Getting Started
Reference Documentation
Advanced JAX Tutorials
Notes
custom_vjp
nondiff_argnums
Developer documentation
API documentation
jax
jax.core
jax.distributed
jax.dlpack
jax.example_libraries
jax.example_libraries.optimizers
jax.example_libraries.stax
jax.experimental.global_device_array
jax.experimental.host_callback
jax.experimental.jet
jax.experimental.loops
jax.experimental.maps
jax.experimental.pjit
jax.experimental.sparse
jax.flatten_util
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.fft
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.betabinom
jax.scipy.stats.cauchy
jax.scipy.stats.chi2
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