NNX – Neural Networks for JAX(github.com)

2 pointsby cgarciae3 lata temu1 comments

https://github.com/cgarciae/nnx

1 comments

NNX is a Neural Networks library for JAX that provides a simple yet powerful module system that adheres to standard Python semantics. Its aim is to combine the robustness of Flax with a simplified, Pythonic API akin to that of PyTorch.