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?
Recently, the
flaxlibrary has unveiled its latest neural network programming framework,nnx. Compared to the originallinen,nnxboasts a plethora of sophisticated features. The architecture has made significant advancements in model persistence and network visualization. Writing networks withnnxhas become more concise and user-friendly. Are there any plans fore3nn_jaxto supportflax.nnx?