Fix neighbor_list_from_ase(compute_distances=True) crashing on periodic structures

[toohardtogetname]

Summary

matgl.ext._alchmtk._compute_distances does

vectors += unit_shifts @ cell

with the unit_shifts tensor exactly as the nvalchemiops neighbor-list
kernel returns it — an integer tensor. Integer × float matmul is not
defined in PyTorch, so the call raises

RuntimeError: expected mat1 and mat2 to have the same dtype, but got: int != float

Because compute_distances=True is the default of
neighbor_list_from_ase, the default call crashes on any structure
with periodic boundary conditions. neighbor_list_from_structure shares
the same sink and is hit even harder: a pymatgen Structure is always
periodic, so its default call crashes on every input. Non-periodic
inputs dodge the path (unit_shifts is None), which is how it survived
smoke tests — and notably, all three internal callers in
matgl.ext.pymatgen / matgl.ext.ase explicitly pass
compute_distances=False, so the default path had zero coverage anywhere.
The one-line fix at _compute_distances repairs all three public entry
points at once.

Who hits this — and who doesn't

Nothing in matgl's own pipelines calls this branch today: model
training/inference and MD go through Structure2Graph, which computes bond
distances elsewhere. Existing results are unaffected, which is also why the
crash went unnoticed. It bites the first direct user of the public API on a
periodic system, e.g.:

• downstream code that reaches for matgl's GPU neighbor list as a building
  block (an MD engine, an adapter, a new-model prototype) — the natural call
  neighbor_list_from_ase(atoms, cutoff) crashes on any crystal (this is how
  we found it, while auditing neighbor lists);
• structure-analysis tooling (coordination numbers, RDF prep, bond
  statistics) — exactly the use case the returned distances exist for, and
  almost always on periodic cells;
• cross-implementation validation, comparing another neighbor-list build
  against matgl's;
• any future matgl-internal refactor that starts using this utility with
  distances — fixing it now removes that landmine before it can reach the
  main pipeline.

Fix

One line — cast the shifts to the position dtype at the point of use:

vectors += unit_shifts.to(positions.dtype) @ cell

Test

tests/ext/test_alchmtk_neighborlist.py: a periodic rocksalt supercell through the
default compute_distances=True path. It asserts not only that the call
survives but that the returned distances match a float64 recomputation
from the returned (positions, shifts, cell) to 1e-4 and respect the
cutoff.

Verified on H200 / CUDA 12.4 / matgl 4.0.3: the test fails before the
patch with the RuntimeError above and passes after it.

While auditing this path

We also cross-checked the neighbor list itself (not part of this PR,
reported for confidence): on a dense triclinic 1,248-atom cell at 5.0 Å,
neighbor_list_from_ase returned an edge set bit-identical (shift
vectors included) to an exact O(N²) minimum-image reference — 25,242
edges, zero missing, zero spurious. The _safe_nl retry-on-overflow logic
works as designed.

🤖 Generated with Claude Code