Add IrDictPolynomial and _ir_dict descriptor variants#271
Open
mariogeiger wants to merge 16 commits intomainfrom
Open
Add IrDictPolynomial and _ir_dict descriptor variants#271mariogeiger wants to merge 16 commits intomainfrom
mariogeiger wants to merge 16 commits intomainfrom
Conversation
Reference documentation for AI coding assistants covering: - cuequivariance: custom group definition (Irrep subclass with @DataClass), segmented tensor products, CG coefficients, descriptors, EquivariantPolynomial - cuequivariance_jax: segmented_polynomial primitive (naive/uniform_1d), RepArray and ir_dict interfaces, NNX layers, indexing, multi-batch axes Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Move SKILL.md files from repo root into the Python package directories
so they ship with pip install. Add __main__.py so users can retrieve
the skill content with:
python -m cuequivariance skill
python -m cuequivariance_jax skill
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
- Fix jax/poly.rst: output ShapeDtypeStruct was (3,3) for a flat size-9 operand, causing a warning. Use (-1,) to infer from descriptor. - Add jax/ir_dict.rst tutorial: split_operand_by_irrep, segmented_polynomial_uniform_1d, gather/scatter indexing, utilities. - Add ir_dict, nnx, and Repeats sections to cuequivariance_jax API docs. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Makes the `python -m <package> skill` command discoverable via `help()` or docstring inspection, which is typically the first thing an AI coding assistant does when exploring an unfamiliar library. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Mirrors what was done for cuequivariance and cuequivariance_jax: - SKILL.md with full API reference for AI coding assistants - __main__.py CLI (`python -m cuequivariance_torch skill`) - Module docstring pointing to the CLI Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Introduce a new ir_dict workflow that decouples from EquivariantPolynomial and IrrepsAndLayout. The key additions: - IrDictPolynomial: a lightweight dataclass pairing a SegmentedPolynomial (already split by irrep) with per-operand-group Irreps metadata. - split_polynomial_by_irreps: standalone helper bridging Irreps and SegmentedPolynomial.split_operand_by_size. - _ir_dict variants for all descriptors (channelwise, fully_connected, full, elementwise, linear, symmetric_contraction, spherical_harmonics). Each descriptor is refactored with a _core function shared between the existing EquivariantPolynomial path and the new IrDictPolynomial path. The NNX layers (nnx.py) no longer import RepArray, EquivariantPolynomial, IrrepsAndLayout, or spherical_harmonics. The MACE and NequIP examples are updated to use the new _ir_dict path. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
The weight irreps changed from per-degree blocks (32x0+80x0+176x0) to a single scalar block (288x0) after the _core extraction. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Skip numpy 1.26 downgrade test on 3.14 since numpy 1.26 does not support Python 3.14. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
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
1 participant
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.
Introduce a new ir_dict workflow that decouples from EquivariantPolynomial and IrrepsAndLayout. The key additions:
The NNX layers (nnx.py) no longer import RepArray, EquivariantPolynomial, IrrepsAndLayout, or spherical_harmonics. The MACE and NequIP examples are updated to use the new _ir_dict path.