Last released Mar 11, 2026
Bijections & normalizing flows with JAX/NNX
Last released Aug 7, 2025
Automatically broadcast inputs by dynamically applying jax.vmap given un-batched argument dimensions.
Last released Oct 21, 2022
Cached lazy evaluation of computational trees
Supported by
