NNX – Neural Networks for JAX (opens in new tab)

(github.com)

2 pointscgarciae3y ago1 comments

1 comments

1 comments · 1 top-level

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.

j / k navigate · click thread line to collapse