Getting Started
jax.debug.print
jax.debug.breakpoint
checkify
Reference Documentation
Advanced JAX Tutorials
Notes
Developer documentation
API documentation
jax
jax.core
jax.debug
jax.distributed
jax.dlpack
jax.example_libraries
jax.example_libraries.optimizers
jax.example_libraries.stax
jax.experimental.checkify
jax.experimental.global_device_array
jax.experimental.host_callback
jax.experimental.jet
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.gennorm
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