Cannot import 'FunctionalFullyConnectedTensorProduct' from 'e3nn_jax'

[ewang26]

Hi! I've installed jax-md from conda but can't seem to import it in a jupyter notebook (or in Python at all). The issue seems to be that the FunctionalFullyConnectedTensorProduct class is in a legacy directory, and it cannot be accessed by the e3nn_jax module. Any help would be appreciated. Sorry if I'm missing something basic here—first time trying to use jax-md!

[MarcBerneman]

It should be fixed now if you build from source. The version installed with pip/conda still hasn't been updated.

git clone https://github.com/google/jax-md
pip install -e jax-md

Otherwise, you might be able to fix this if you use an older version of JAX.