Enquiring About e3nn_jax Support for the New flax.nnx Framework

[kYangLi]

Recently, the flax library has unveiled its latest neural network programming framework, nnx. Compared to the original linen, nnx boasts a plethora of sophisticated features. The architecture has made significant advancements in model persistence and network visualization. Writing networks with nnx has become more concise and user-friendly. Are there any plans for e3nn_jax to support flax.nnx?

[ameya98]

Hmm I could give it a shot sometime next week!

[chen-zz20]

I am looking forward to the new version of the program. Thank you!

[ameya98]

Sorry for not looking at this for a while, but it looks pretty cool. My plan is to implement a Linear layer in nnx by the end of the month.