Numerics and Systems Lab in JAX: 21 notebooks executed end to end in CI, plus jax_nsl, a 275-test library covering numerically stable primitives, autodiff tools (Hessian-vector products, implicit differentiation), transforms, preconditioned linear algebra, training code and multi-device parallelism.
machine-learning research deep-learning optimization automatic-differentiation parallel-computing transformers bayesian-methods scientific-computing neural-networks high-performance-computing differential-equations gpu-computing pmap distributed-training jax mixed-precision physics-informed pjit tpu-computing
-
Updated
Nov 1, 2026 - Jupyter Notebook