Skip to content

2023-05-02
Compare
Choose a tag to compare
@mariogeiger mariogeiger released this 02 May 19:46
· 201 commits to main since this release
0.17.4
b793251

Highlight

Finally implemented the eSCN optimization. That is the optimization of the tensor product with spherical harmonics followed by a linear mix.

old:

sh = e3nn.spherical_harmonics(vector, range(lmax), True)
features = e3nn.tensor_product(features, sh)
linear = e3nn.flax.Linear(irreps_out)
features = linear.apply(w, features)

new: (equivalent but faster)

from e3nn_jax.experimental.linear_shtp import LinearSHTP
conv = LinearSHTP(irreps_out)
features = conv.apply(w, features, vector)

Example of usage:

ChangeLog

Added

  • LinearSHTP module implementing the optimized linear mixing of inputs tensor product with spherical harmonics
  • D_from_axis_angle
  • to_s2grid: quadrature="gausslegendre" by default
  • soft_odd activation function for odd scalars
  • more support of arrays implicitely converted into IrrepsArray as scalars (i.e. added few IrrepsArray.as_irreps_array)

Changed

  • scalar_activation simpler to use with default activation functions (a bit like gate)
Assets 2
Loading