Enquiring About `e3nn_jax` Support for the New `flax.nnx` Framework

e3nn / e3nn-jax

jax library for E3 Equivariant Neural Networks

Apache License 2.0

181 stars 18 forks source link

Enquiring About `e3nn_jax` Support for the New `flax.nnx` Framework #80

Open kYangLi opened 1 month ago

kYangLi commented 1 month ago

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 commented 1 month ago

Hmm I could give it a shot sometime next week!

chen-zz20 commented 1 month ago

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