Releases: e3nn/e3nn-jax
Releases · e3nn/e3nn-jax
0.20.8
What's Changed
- Fix dialect for JAX mhlo -> stablehlo. by @ameya98 in #85
- Add support for signals on SO3. by @ameya98 in #84
- Fix SO3 sampling. by @ameya98 in #86
- Don't pre-multiply by angle measure. by @ameya98 in #87
- Add tests. by @ameya98 in #88
- Minor fixes for SO3Signal by @ameya98 in #89
- Fix multiplication. by @ameya98 in #90
- Add argmax for signals on SO3 by @ameya98 in #91
- Fix formatting. by @ameya98 in #92
- Add apply() for SO3Signal. by @ameya98 in #94
- Add division and negation for SO3Signal. by @ameya98 in #95
- Add static fields to e3nn_jax.equinox.Linear by @teddykoker in #96
- Add equinox modules to docs by @teddykoker in #97
New Contributors
- @teddykoker made their first contribution in #96
Full Changelog: 0.20.7...0.20.8
0.20.7
2024-01-26
Added
e3nn.wherefunction- Add optional
maskargument ine3nn.flax.BatchNorm
Changed
- replace
jnp.ndarraybyjax.Array
2024-01-05
Added
e3nn.onesande3nn.ones_likefunctionse3nn.equinoxsubmodule
Fixed
- python 3.9 compatibility
Thanks to @ameya98, @SauravMaheshkar and @pabloferz
2023-12-24
2023-11-17
Added
e3nn.flax.BatchNorme3nn.scatter_mean- Add
e3nn.utils.vmapalso directly toe3nnmodule:e3nn.vmap
2023-09-25
Added
with_biasargument toe3nn.haiku.MultiLayerPerceptronande3nn.flax.MultiLayerPerceptron
Fixed
- Improve compilation speed and stability of
s2gridfor largelmax(useis_normalized=Trueinlpmn_values)
2023-09-13
Changelog
Changed
- Add back the optimizations with the lazy
._chunksthat was removed in 0.19.0
2023-09-09
Highlight
tl;dr Mostly fix the issue #38
In version 0.19.0, I removed the lazy _list attribute of IrrepsArray to fix the issues from tree_util, grad and vmap.
In this version (0.20.0) I found a way to put back that lazy attribute, now called _chunks, in a way that does not interfere with tree_util, grad and vmap. _chunks is tropped when using tree_util, grad and vmap unless you use e3nn.vmap.
ChangeLog
Added
e3nn.Irreps.mul_gcde3nn.IrrepsArray.extend_with_zerosto extend an array with zeros, can be useful for residual connections
Changed
- rewrite
e3nn.tensor_squareto be simpler (and faster?) - use
jax.scipy.special.lpmn_valuesto implemente3nn.legendre. Faster on GPU and supports reverse-mode differentiation. - [BREAKING] Change the output format of
e3nn.legendre!
Fixed
- Add back a lazy
._chunksine3nn.IrrepsArrayto fix issue #38
2023-06-24
Changelog
Fixed
- Fix missing support for zero flags in
e3nn.elementwise_tensor_product