Conversation
Replace raw e3nn.o3.Linear with SevenNet's IrrepsLinear wrapper so the readout supports lazy instantiation, FlashTP/cueq backend conversion, and modality-aware one-hot concatenation. zero_init is deferred to instantiate() to remain compatible with lazy mode. https://claude.ai/code/session_018xjQvBPU2gzbRKMut4EgPa
Weight tensor moved from first_linear.weight to first_linear.linear.weight after refactoring LatentChargeReadout to wrap IrrepsLinear instead of raw e3nn.o3.Linear. https://claude.ai/code/session_018xjQvBPU2gzbRKMut4EgPa
Remove set_num_modalities, _is_batch_data property, and the saved_input restore logic from forward(). Mirrors the LES branch's fix commit (e1af8d3): modality dependence already flows through the upstream conv stack's modality-aware features, so the q-readout does not need to concat the one-hot itself. https://claude.ai/code/session_018xjQvBPU2gzbRKMut4EgPa
test_les_legacy.py:
- test_is_batch_data_propagates: verify set_is_batch_data reaches
EdgePreprocess, LatentEwaldSum, and ForceStressOutput. Guards against
regressions from removing the _is_batch_data property on
LatentChargeReadout.
- test_state_dict_round_trip: save/load round-trip reproduces output
bit-for-bit on a freshly built model.
- test_lr_energy_is_scalar: non-batch LR_ENERGY must be shape () so
AddLREnergy does not silently broadcast onto scalar SR_ENERGY.
- test_lr_changes_force / test_lr_changes_stress: with shared SR weights,
non-zero charges must produce non-trivial force/stress contributions.
Catches a silently broken LR gradient path.
- test_strain_leaf_created_during_forward: also assert is_leaf so
d(E)/d(_strain) is well-defined.
test_les_finetune.py:
- test_is_batch_data_propagates: same regression guard on the omni model.
- test_state_dict_round_trip: save/load through build_model_with_les.
- test_lr_energy_zero_at_init: also assert shape == () in non-batch mode.
https://claude.ai/code/session_018xjQvBPU2gzbRKMut4EgPa
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
2 participants
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.
No description provided.