.gitignore
.readthedocs.yaml
CHANGELOG.md
LICENSE
README.md
pyproject.toml
.github/workflows/lint.yml
.github/workflows/release.yml
.github/workflows/tests.yml
docs/BENCHMARK.rst
docs/README.md
docs/conf.py
docs/favicon.ico
docs/index.rst
docs/requirements.txt
docs/api/equinox.rst
docs/api/extra.rst
docs/api/flax.rst
docs/api/haiku.rst
docs/api/index.rst
docs/api/irreps.rst
docs/api/irreps_array.rst
docs/api/legacy.rst
docs/api/math.rst
docs/api/nn.rst
docs/api/operations.rst
docs/api/radial.rst
docs/api/s2.rst
docs/api/tensor_products.rst
docs/api/utils.rst
docs/tuto/graph.png
docs/tuto/index.rst
docs/tuto/mp-169.png
docs/tuto/nequip.png
docs/tuto/nequip.rst
e3nn_jax/__init__.py
e3nn_jax/equinox.py
e3nn_jax/flax.py
e3nn_jax/haiku.py
e3nn_jax/legacy.py
e3nn_jax/perm.py
e3nn_jax/utils.py
e3nn_jax.egg-info/PKG-INFO
e3nn_jax.egg-info/SOURCES.txt
e3nn_jax.egg-info/dependency_links.txt
e3nn_jax.egg-info/requires.txt
e3nn_jax.egg-info/top_level.txt
e3nn_jax/_src/J.py
e3nn_jax/_src/__init__.py
e3nn_jax/_src/activation.py
e3nn_jax/_src/basic.py
e3nn_jax/_src/config.py
e3nn_jax/_src/dropout_haiku.py
e3nn_jax/_src/fc_tp_haiku.py
e3nn_jax/_src/gate.py
e3nn_jax/_src/grad.py
e3nn_jax/_src/irreps.py
e3nn_jax/_src/irreps_array.py
e3nn_jax/_src/linear.py
e3nn_jax/_src/linear_equinox.py
e3nn_jax/_src/linear_flax.py
e3nn_jax/_src/linear_haiku.py
e3nn_jax/_src/mlp_flax.py
e3nn_jax/_src/mlp_haiku.py
e3nn_jax/_src/radial.py
e3nn_jax/_src/radius_graph.py
e3nn_jax/_src/reduced_tensor_product.py
e3nn_jax/_src/rotation.py
e3nn_jax/_src/s2grid.py
e3nn_jax/_src/scatter.py
e3nn_jax/_src/so3.py
e3nn_jax/_src/so3grid.py
e3nn_jax/_src/su2.py
e3nn_jax/_src/symmetric_tensor_product_haiku.py
e3nn_jax/_src/tensor_product_with_spherical_harmonics.py
e3nn_jax/_src/tensor_products.py
e3nn_jax/_src/batchnorm/__init__.py
e3nn_jax/_src/batchnorm/bn.py
e3nn_jax/_src/batchnorm/bn_flax.py
e3nn_jax/_src/batchnorm/bn_haiku.py
e3nn_jax/_src/legacy/__init__.py
e3nn_jax/_src/legacy/core_tensor_product.py
e3nn_jax/_src/legacy/functional_fully_connected_tensor_product.py
e3nn_jax/_src/spherical_harmonics/__init__.py
e3nn_jax/_src/spherical_harmonics/legendre.py
e3nn_jax/_src/spherical_harmonics/recursive.py
e3nn_jax/_src/utils/__init__.py
e3nn_jax/_src/utils/decorators.py
e3nn_jax/_src/utils/dtype.py
e3nn_jax/_src/utils/einsum.py
e3nn_jax/_src/utils/jit.py
e3nn_jax/_src/utils/math_numpy.py
e3nn_jax/_src/utils/optimize_jaxpr.py
e3nn_jax/_src/utils/sum_tensors.py
e3nn_jax/_src/utils/sympy.py
e3nn_jax/_src/utils/test.py
e3nn_jax/_src/utils/vmap.py
e3nn_jax/experimental/__init__.py
e3nn_jax/experimental/linear_shtp.py
e3nn_jax/experimental/point_convolution.py
e3nn_jax/experimental/transformer.py
e3nn_jax/experimental/voxel_convolution.py
e3nn_jax/experimental/voxel_pooling.py
e3nn_jax/irrep/__init__.py
examples/__init__.py
examples/irreps_array.py
examples/plot_spherical_harmonics.py
examples/tensor_product_benchmark.py
examples/tetris_point.py
examples/tetris_point_test.py
examples/tetris_voxel.py
examples/tetris_voxel_test.py
tests/conftest.py
tests/noxfile.py
tests/_src/activation_test.py
tests/_src/basic_test.py
tests/_src/config_test.py
tests/_src/dropout_haiku_test.py
tests/_src/gate_test.py
tests/_src/grad_test.py
tests/_src/irreps_array_test.py
tests/_src/irreps_test.py
tests/_src/linear_equinox_test.py
tests/_src/linear_flax_test.py
tests/_src/linear_haiku_test.py
tests/_src/perm_test.py
tests/_src/radial_test.py
tests/_src/radius_graph_test.py
tests/_src/reduced_tensor_product_test.py
tests/_src/rotation_test.py
tests/_src/s2grid_test.py
tests/_src/scatter_test.py
tests/_src/so3_test.py
tests/_src/so3grid_test.py
tests/_src/symmetric_tensor_product_haiku_test.py
tests/_src/tensor_products_test.py
tests/_src/batchnorm/bn_flax_test.py
tests/_src/batchnorm/bn_haiku_test.py
tests/_src/legacy/core_tensor_product_test.py
tests/_src/spherical_harmonics/sh_test.py
tests/_src/utils/einsum_test.py
tests/_src/utils/jit_test.py
tests/experimental/linear_shtp_test.py
tests/experimental/point_convolution_test.py
tests/experimental/transformer_test.py
tests/experimental/voxel_convolution_test.py
tests/experimental/voxel_pooling_test.py