numpy
e3nn_jax
ordered_set
flax
einops
wandb

[dev]

[test]
pytest
