Conversation
``get_int_dtype()`` (``TORCHREF_DTYPE_INT``, int32 by default) now covers every integer tensor the package allocates or casts, replacing 205 literal ``torch.long`` / ``int64`` / ``int32`` / ``int8`` dtypes that mostly carried an "indexing requires long" justification. That has not been true since torch 2.0: bracket indexing, ``index_select`` and ``index_add_`` accept int32, and those are the consumers of nearly every converted site. The literal int64 that remains is what torch or the arithmetic forces, and each site's marker names the constraint: - ``scatter_add`` / ``gather`` indices, which need int64 on torch < 2.8; the declared ``torch>=2.4`` floor still covers those releases. - Packed keys -- ``min(i, j) * max_idx + max(i, j)`` pair hashes, the composite HKL sort key and the clustering keys -- which overflow int32, and which ``searchsorted`` needs in the same dtype as the table. - TorchMD-Net's ``Z`` and ``batch`` tensors, an external library contract. Count accumulators that ``scatter_add_`` or ``index_add_`` a ``ones_like`` of their index take that index's dtype instead of a literal, so the source and self dtypes match whatever the index is. ``AtomGraph.implicit_h_count`` brings ``bincount``'s int64 to the configured dtype once instead of widening the template count. The riding-hydrogen pair hash widens its operands to int64 explicitly, since the candidate indices it packs are no longer int64 by construction. The 18 markers the dtype guard flagged on dev sat on the closing line of a multi-line call, where the checker never looks; converting those sites removes the markers altogether. AGENTS.md states the convention and its exceptions, and the anchor-selection test compares against the configured dtype rather than int64. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01KuUn93S8goPxwJLBWjjhht
``dest[index] = source`` checks that the two dtypes match, unlike slice assignment, so an int64 ``arange`` or ``unique`` inverse can no longer be written into a table that now carries the configured int dtype. The Friedel-mate maps, the grid-cell lookup, the spherical-harmonic row map and the atom remap used by hydrogen generation build their sources in the table's dtype instead. The Legendre shell kernel TORCH_CHECKs int64 shell labels and offsets, so those two tensors keep int64 and their markers say why. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01KuUn93S8goPxwJLBWjjhht
Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01KuUn93S8goPxwJLBWjjhht
```_reindex_per_reflection``` clones the new HKL table, which carries the configured int dtype, and writes the stored anomalous HKL rows into it. Under ``TORCHREF_DTYPE_INT=int64`` those rows were int32, and indexed assignment rejects the mismatch; this predates the dtype conversion and showed up in the int64 run of the reflection-data tests. Cast the rows to the table's dtype. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01KuUn93S8goPxwJLBWjjhht
Resolve the conflicts from dev moving the shell helpers out of sigma_a.py: the segment layout in _shells.py takes the configured int dtype, the same conversion this branch had applied to it before the move, and the changelog keeps both sides. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01KuUn93S8goPxwJLBWjjhht
Lines that grew past 88 columns when their literal became get_int_dtype() are wrapped, with each remaining dtype-ok marker on the line above its literal so the dtype guard still sees it. The configured-int imports sit in their first-party groups, the canonical HKL sort key casts once before slicing, and two docstring lines are rewrapped. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01KuUn93S8goPxwJLBWjjhht
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: b4d6f91812
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| ki = torch.remainder(k, Ny) | ||
| li = torch.remainder(l, Nz) | ||
| lin = (hi * (Ny * Nz) + ki * Nz + li).to(torch.int64) # (N,) # dtype-ok: flat grid index (lin) for scatter/gather; requires int64 | ||
| lin = (hi * (Ny * Nz) + ki * Nz + li).to(get_int_dtype()) # (N,) |
There was a problem hiding this comment.
When the default integer configuration is int32 and Nx * Ny * Nz exceeds 2**31 - 1, this packed grid-index expression overflows before the final .to(get_int_dtype()); index_add then receives negative or aliased indices, so large reciprocal grids either fail or place structure factors in the wrong voxels. Compute lin and lin_sym in int64, which is precisely the packed-key exception to the configured-dtype policy.
AGENTS.md reference: AGENTS.md:L54-L59
Useful? React with 👍 / 👎.
| ki = torch.remainder(all_hkl[:, 1], Ny) | ||
| li = torch.remainder(all_hkl[:, 2], Nz) | ||
| return (hi * (Ny * Nz) + ki * Nz + li).to(torch.int64) # dtype-ok: flat HKL grid index; int64 avoids overflow, used for indexing | ||
| return (hi * (Ny * Nz) + ki * Nz + li).to(get_int_dtype()) |
There was a problem hiding this comment.
For a reciprocal grid containing more than 2**31 - 1 elements under the default int32 configuration, hi * (Ny * Nz) + ki * Nz + li wraps before the result is used by gather_with_index_add, causing out-of-range access or extraction from aliased voxels. This flattened index is a packed key and must remain int64 rather than following the general configured-index dtype.
AGENTS.md reference: AGENTS.md:L54-L59
Useful? React with 👍 / 👎.
| S, N = G.shape | ||
| coeff = obs.coeff.to(device=device, dtype=cplx) | ||
| h_R_int = cand.h_R.round().to(torch.int64) # dtype-ok: Miller indices are integers | ||
| h_R_int = cand.h_R.round().to(get_int_dtype()) |
There was a problem hiding this comment.
When the translation grid has more than 2**31 - 1 points, converting h_R_int to the default int32 makes the later flattened-grid expression overflow before W.index_add_, producing invalid or colliding translation bins. Retain int64 for this tensor, or cast before forming flat, because this is packed-index arithmetic rather than an ordinary index tensor.
AGENTS.md reference: AGENTS.md:L54-L59
Useful? React with 👍 / 👎.
Codex Review SummaryThis comment shows the latest Codex review activity on this pull request.
ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings. |
This branch has not been deployed
2 participants
Summary
Integer and index tensors throughout
torchref/now take the configured int dtype (get_int_dtype(),TORCHREF_DTYPE_INT, int32 by default) instead of a hardcodedtorch.int64/torch.long. This resolves the dtype-guard category of failures seen on #71 by following the policy rather than marking every literal.Changes
get_int_dtype(); the one test that pinned an int64 dtype now pins the configured one.# dtype-ok:marker naming the constraint:scatter/gatherfamily indices: int64 required on torch < 2.8 (the repo floor is 2.4)index_copy_,index_fill_,one_hot: int64 alwaysi * n + jthat overflow int32TORCH_CHECKs int64 labels and offsetsZas a LongTensor)dest[idx] = srcrequires matching dtypes.ReflectionDataanomalous reindexing casts gathered values to the destination dtype; it raised underTORCHREF_DTYPE_INT=int64before.AGENTS.md§2.1 states the policy and the exception list;docs/changelog.rstcarries the user-facing line.dev(5f32d32) is merged in; the segment layout that moved into_shells.pytakes the configured int dtype as well.Verification
pytest tests/ -m "not gpu and not slow"on CPU: 2648 passed, 109 skipped, 7 failed. Six fail identically on unmodifieddev: fivetest_af_trajectory_matches_referencecases, whose R-work now ends 0.026–0.042 below the recorded reference, andtest_copy_after_refinement_setup_is_broken_for_every_representation. The seventh,test_dataset_scaler.py::test_free_and_validation_changes_do_not_affect_fit, is a pre-existing hash-seed flake: on unmodifieddevit fails for 5 of 8 values ofPYTHONHASHSEED(0, 1, 2, 5, 7) and on this branch for 2 of 8, so some iteration order in the scaling path depends on string hashing. Not touched here.tests/unit/test_dtype_conformance.pypasses: every remaining int literal is justified.TORCHREF_DTYPE_INT=int64,pytest tests/unit -m "not gpu and not slow": 2299 passed, 85 skipped, 0 failed.dev. black and isort change nothing on the lines this PR adds; both want to reformat much of the repository independently of this PR, so whole-file checks stay noisy.🤖 Generated with Claude Code
https://claude.ai/code/session_01KuUn93S8goPxwJLBWjjhht
Generated by Claude Code