Tutorials
Advanced JAX Tutorials
Notes
Developer documentation
API documentation
jax
jax.core
jax.dlpack
jax.experimental
jax.experimental.host_callback
jax.experimental.loops
jax.experimental.optimizers
jax.experimental.stax
jax.image
jax.lax
jax.lax.linalg
jax.nn
jax.nn.initializers
jax.numpy
jax.numpy.fft
jax.numpy.linalg
jax.ops
jax.profiler
jax.random
jax.scipy.linalg
jax.scipy.ndimage
jax.scipy.optimize
jax.scipy.signal
jax.scipy.sparse.linalg
jax.scipy.special
jax.scipy.stats.bernoulli
jax.scipy.stats.beta
jax.scipy.stats.cauchy
jax.scipy.stats.dirichlet
jax.scipy.stats.expon
jax.scipy.stats.gamma
jax.scipy.stats.geom
jax.scipy.stats.laplace
jax.scipy.stats.logistic
jax.scipy.stats.multivariate_normal
jax.scipy.stats.norm
jax.scipy.stats.pareto
jax.scipy.stats.poisson
jax.scipy.stats.t
jax.scipy.stats.uniform
jax.tree_util