Fix state.to() not propagating dtype/device to constraints & expose angle_tolerance in symmetry pipeline#527
Merged
CompRhys merged 10 commits intoTorchSim:mainfrom Apr 2, 2026
Conversation
thomasloux
approved these changes
Apr 1, 2026
added 3 commits
April 1, 2026 15:19
Member
|
just the linting and then good to merge? |
Contributor
Author
|
Hopefully good to go -- I've fixed the last remaining linting error that was introduced by this PR. The others are pre-existing I believe. |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
state.to(dtype=...)did not update constraint tensors._state_to_deviceonly converted top-levelSimStatetensor attributes (positions,cell,masses) but left constraint objects untouched. This meant internal tensors likeFixSymmetry.rotationsorFixCom.comsstayed on the original dtype/device, causing silent dtype mismatches.state.to(dtype=...)did not update sub-class tensor attributes likeforces,velocities, ....angle_tolerancecould not be passed to moyopy for symmetry detection.MoyoDatasetsupports anangle_toleranceparameter, but it wasn't exposed anywhere in the TorchSim symmetry pipeline.Changes
Constraint dtype/device propagation
torch_sim/constraints.py,torch_sim/state.pyConstraint.to(device, dtype)as a base no-op, with concrete overrides on each subclass:to()movesAtomConstraintatom_idxSystemConstraintsystem_idxFixComsystem_idx+ cachedcomstensorFixSymmetryrotations(float→dtype),symm_maps(int→device only),reference_cellsto()follows the same explicit-constructor pattern used byreindex,merge, andselect_constraint._state_to_devicenow calls.to()on each attached constraint._state_to_devicedtype handlingtorch_sim/state.py.to(device=device, dtype=dtype), integer/bool tensors get.to(device=device)only.positions/masses/cellfor dtype conversion, so subclass attributes likeforcesandmomentaare now handled automatically.Expose
angle_tolerancetorch_sim/symmetrize.py,torch_sim/constraints.pyangle_tolerance: float | None = Noneparameter to:_moyo_datasetget_symmetry_datasetsprep_symmetry_refine_symmetry_implrefine_symmetryrefine_and_prep_symmetryFixSymmetry.from_state().MoyoDataset(…, angle_tolerance=angle_tolerance).Checklist
Before a pull request can be merged, the following items must be checked: