google/neural-tangents
A JAX-based library for defining finite and infinite-width neural networks with theoretical guarantees from neural network kernel theory.

Neural Tangents enables researchers to specify complex, hierarchical neural networks of both finite and infinite width through a high-level API. It leverages neural network kernel theory to provide insights into training dynamics and correspondence with Gaussian processes. The library is built on JAX for automatic differentiation and scalable computation, supporting gradient descent analysis and Bayesian inference in deep networks.