Skip to content

2023-09-09

Compare
Choose a tag to compare
@mariogeiger mariogeiger released this 09 Sep 18:58
· 101 commits to main since this release

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_gcd
  • e3nn.IrrepsArray.extend_with_zeros to extend an array with zeros, can be useful for residual connections

Changed

  • rewrite e3nn.tensor_square to be simpler (and faster?)
  • use jax.scipy.special.lpmn_values to implement e3nn.legendre. Faster on GPU and supports reverse-mode differentiation.
  • [BREAKING] Change the output format of e3nn.legendre!

Fixed

  • Add back a lazy ._chunks in e3nn.IrrepsArray to fix issue #38