SphericalSignal in e3nn-JAX

A small demonstration of the SphericalSignal API in e3nn-JAX.

This is the JAX version of https://gist.github.com/blondegeek/e37acc6091a8b9787fad1c7536617d44.

Installation

With pip:

python3 -m venv .venv
source .venv/bin/activate
pip install -r requirements.txt

With uv:

uv venv
source .venv/bin/activate
uv pip install -r requirements.txt