ASCL Code Record
[ascl:2504.026]
Flax: Neural network library for JAX
Heek, Jonathan; Levskaya, Anselm; Oliver, Avital; Ritter, Marvin; Rondepierre, Bertrand; Steiner, Andreas; van Zee, Marc
Flax provides a flexible end-to-end user experience for JAX users; its NNX is a simplified API that creates, inspects, debugs, and analyzes neural networks in JAX. It has first class support for Python reference semantics, enabling users to express their models using regular Python objects. Flax NNX is an evolution of the previous Flax Linen API.
- Code site:
- https://github.com/google/flax
- Used in:
- https://ui.adsabs.harvard.edu/abs/2025ApJ...980...66R https://ui.adsabs.harvard.edu/abs/2025JCAP...01..095L
- Bibcode:
- 2025ascl.soft04026H
Views: 14
