allscaip: fix radius-graph dtype mismatches under float64 inference

[OwenPriceSkelly]

What does this PR do?

Fixes two dtype crashes in AllScAIP's radius-graph construction when running float64 inference (InferenceSettings(base_precision_dtype=torch.float64)):

1. biknn_radius_graph builds the PBC image-offset tensors (image_id) with torch.get_default_dtype() (float32), while the fp64 batch's cell is double:

   File "src/fairchem/core/models/allscaip/utils/allscaip_radius_graph.py", line 529, in build_radius_graph
       src_pos = pos[:, None] + torch.mm(image_id, cell)[None, :]
   RuntimeError: expected mat1 and mat2 to have the same dtype, but got: float != double
   

2. Once past that, the padded output buffers (padded_disp, src_env, dst_env) are also created at the default dtype and then index-put with double disp/env:

   RuntimeError: Index put requires the source and destination dtypes match, got Float for the destination and Double for the source.
   

Fix: build image_id with data.cell.dtype and the padded buffers with the dtypes of the tensors scattered into them. No behavior change at the default float32 (the new dtype expressions resolve to float32 there).

Found running allscaip-md-{conserving,direct}-all-omol at fp64 through pretrained_mlip.get_predict_unit + FAIRChemCalculator on ALCF Aurora (torch 2.13, xpu), but the bug is device-independent — the added test reproduces both crashes on CPU at upstream main.

Test

tests/core/models/allscaip/test_radius_graph_dtype.py — parametrized over float32/float64 × PBC/non-PBC; the float64 cases fail on main (first crash above; after fixing only image_id, the second) and pass with this change. test_allscaip_forward.py still passes; ruff check/ruff format clean.

🤖 Generated with Claude Code

[meta-cla]

Hi @OwenPriceSkelly!

Thank you for your pull request and welcome to our community.

Action Required

In order to merge any pull request (code, docs, etc.), we require contributors to sign our Contributor License Agreement, and we don't seem to have one on file for you.

Process

In order for us to review and merge your suggested changes, please sign at https://code.facebook.com/cla. If you are contributing on behalf of someone else (eg your employer), the individual CLA may not be sufficient and your employer may need to sign the corporate CLA.

Once the CLA is signed, our tooling will perform checks and validations. Afterwards, the pull request will be tagged with CLA signed. The tagging process may take up to 1 hour after signing. Please give it that time before contacting us about it.

If you have received this in error or have any questions, please contact us at cla@meta.com. Thanks!

[meta-cla]

Thank you for signing our Contributor License Agreement. We can now accept your code for this (and any) Meta Open Source project. Thanks!

[OwenPriceSkelly]

Scope note from further testing: this PR fixes the radius-graph construction path (both crashes above reproduce and pass in the added test), but it is not the whole float64 story for AllScAIP. Running a full tiny-model forward at fp64 on CPU (mirroring test_allscaip_forward.py's config, model.to(torch.float64), double batch) gets past the graph with this patch and then hits:

RuntimeError: mat1 and mat2 must have the same dtype, but got Float and Double

at the energy head's FFN — the backbone hard-casts its node representations to float32 before handing them to heads (AllScAIP.py: "node_reps": neighbor_reps[:, 0].to(torch.float32)), so double head weights mismatch. Since that cast looks deliberate (fp32 heads regardless of base precision?), I've left it out of this PR rather than guess at the intended design — but if InferenceSettings(base_precision_dtype=torch.float64) is meant to be supported for AllScAIP end-to-end, that seam (and the explicit torch.float32 masks in utils/data_preprocess.py) would need a decision from the model owners. Happy to extend this PR or file it separately, whichever you prefer.