Refactor core components, update CI/infrastructure, and improve test stability - #102
Refactor core components, update CI/infrastructure, and improve test stability #102muhrin wants to merge 10 commits into
Conversation
- Update radial, s2grid, and scatter modules - Refactor tensor_product_with_spherical_harmonics - Improve utility decorators, vmap, and test helpers - Update related benchmarks and tests
- Update requires-python to >=3.10 and add classifiers for 3.10-3.14 - Set black target-version to py310 to ensure broader compatibility - Configure flake8 to read from pyproject.toml via flake8-pyproject
- Increase sample size in tests/_src/gate_test.py to 2048 to reduce variance. - Loosen assertion thresholds in tests/_src/linear_haiku_test.py and tests/_src/legacy/core_tensor_product_test.py to accommodate increased precision sensitivity. - Refactor internal type handling and imports in e3nn_jax/_src/scatter.py, e3nn_jax/_src/spherical_harmonics/legendre.py, e3nn_jax/_src/utils/test.py, tests/_src/grad_test.py, and tests/_src/irreps_array_test.py for consistency.
|
Hey @ameya98 , Martin Uhrin here, I'm friend with Tess and Mario. If you get a chance to review this that would be great. For our codebase it would be super useful to be able to get this released relatively soon. Thanks! |
|
Hey thanks for the PR. Yes separate PR for the CI part would be helpful. Can you also toss the test changes for x64 and formatting into that. Keep the refactor part in this PR. Will try to give a review later this week. |
|
Happily, I'll try and repackage it tomorrow. |
|
cool to see new PR here ;) |
|
Hey! Great to read you :) Indeed, our whole code stack depends on this library and, let's say, it needed a bit of tender loving care to bring things up to date with changes in python/JAX. So you can count on us to keep it alive for at least some time :) |
3 participants
Some CI work, updated for python 3.10-3.14. No API breaking changes. I've also fixed a bug to do with forward mode AD and calculation of Jacobians.
If you'd prefer this spit into multiple PRs, let me know.
Key Changes
- Core & Refactoring:
- Updated core tensor product and utility operations.
- Simplified jit_code to return StableHLO MLIR directly.
- Refactored IrrepsArray Jacobian calculation, Irreps, and Spherical Harmonics.
- CI & Infrastructure:
- Dropped support for Python 3.9 and added support up to 3.14.
- Updated Python version support and improved linting configuration.
- Updated GitHub CI workflows.
- Testing:
- Fixed statistical flakiness in tests when JAX x64 is enabled.
- Adjusted tolerances and configurations for spherical harmonics tests.
- Included minor formatting and test modifications.