All modules for which code is available
- DeviceArray
- jax._src.api
- jax._src.custom_derivatives
- jax._src.dtypes
- jax._src.errors
- jax._src.flatten_util
- jax._src.image.scale
- jax._src.lax.control_flow
- jax._src.lax.convolution
- jax._src.lax.fft
- jax._src.lax.lax
- jax._src.lax.linalg
- jax._src.lax.other
- jax._src.lax.parallel
- jax._src.lax.qdwh
- jax._src.lax.slicing
- jax._src.lax.windowed_reductions
- jax._src.lib.xla_bridge
- jax._src.nn.functions
- jax._src.nn.initializers
- jax._src.numpy.fft
- jax._src.numpy.lax_numpy
- jax._src.numpy.linalg
- jax._src.numpy.polynomial
- jax._src.numpy.vectorize
- jax._src.ops.scatter
- jax._src.profiler
- jax._src.random
- jax._src.scipy.fft
- jax._src.scipy.linalg
- jax._src.scipy.ndimage
- jax._src.scipy.optimize.minimize
- jax._src.scipy.signal
- jax._src.scipy.sparse.linalg
- jax._src.scipy.special
- jax._src.scipy.stats.bernoulli
- jax._src.scipy.stats.beta
- jax._src.scipy.stats.betabinom
- jax._src.scipy.stats.cauchy
- jax._src.scipy.stats.chi2
- jax._src.scipy.stats.dirichlet
- jax._src.scipy.stats.expon
- jax._src.scipy.stats.gamma
- jax._src.scipy.stats.geom
- jax._src.scipy.stats.laplace
- jax._src.scipy.stats.logistic
- jax._src.scipy.stats.multivariate_normal
- jax._src.scipy.stats.norm
- jax._src.scipy.stats.pareto
- jax._src.scipy.stats.poisson
- jax._src.scipy.stats.t
- jax._src.scipy.stats.uniform
- jax._src.third_party.numpy.linalg
- jax._src.tree_util
- jax.core
- jax.example_libraries.optimizers
- jax.example_libraries.stax
- jax.experimental.ann
- jax.experimental.host_callback
- jax.experimental.loops
- jax.experimental.maps
- jax.experimental.pjit
- jax.experimental.sparse.bcoo
- jax.experimental.sparse.transform
- jax.experimental.x64_context
- jaxlib.xla_extension
- numpy.core.getlimits