RFD3 error with inference_sampler.allow_realignment option

[ben-jsmith]

I'm getting this error in RFdiffusion3 when running motif scaffolding with allow_realignment=True

foundry/utils/alignment.py:60 in weighted_rigid_align                          │
│                                                                              │
│   57 │   # Computation of the covariance matrix                              │
│   58 │   C = torch.einsum("bji,bjk->bik", w_resolved[..., None] * X_gt_resol │
│   59 │                                                                       │
│ ❱ 60 │   U, S, V = torch.linalg.svd(C)                                       │
│   61 │                                                                       │
│   62 │   R = U @ V                                                           │
│   63 │   B, _, _ = X_L.shape                                                 │

RuntimeError: "svd_cuda_gesvdjBatched" not implemented for 'BFloat16'

I'm wondering if this is a bug or a mistake in my command

rfd3 design \
        out_dir=test \
        inputs=input.json \
        dump_trajectories=True \
        prevalidate_inputs=True \
        inference_sampler.allow_realignment=True \
        n_batches=1 \
        diffusion_batch_size=1

with this example json file used as input:

 {
     "test": {
         "input": "7v11.pdb",
         "ligand": "OQO",
         "unindex": "A431,A572-573",
         "length": "200",
         "select_fixed_atoms": {
             "A431": "TIP",
             "A572": "BKBN",
            "A573": "BKBN"
        }
    }
}

The error does not occur when doing unconditional generation with allow_realignment=True set or when doing motif scaffolding without the argument.

[r-krishna]

Oh interesting. That almost certainly looks like a bug to me. Also tagging @Ubiquinone-dot for his thoughts. There might be somewhere where it is not upcasted in the motif scaffolding case.

[r-krishna]

We will work on getting it resolved ASAP!

[Ubiquinone-dot]

Generally it's not needed to have allow_realignment=True, we found fewer weird interactions with motif scaffolding when it's off - but the code should work so will take a closer look soon!