All modules for which code is available
- jax._src.ad_checkpoint
- jax._src.api
- jax._src.array
- jax._src.callback
- jax._src.checkify
- jax._src.compilation_cache
- jax._src.compiler
- jax._src.config
- jax._src.core
- jax._src.custom_derivatives
- jax._src.debugger.core
- jax._src.debugging
- jax._src.distributed
- jax._src.dlpack
- jax._src.dtypes
- jax._src.environment_info
- jax._src.extend.random
- jax._src.flatten_util
- jax._src.image.scale
- jax._src.lax.ann
- jax._src.lax.control_flow.conditionals
- jax._src.lax.control_flow.loops
- jax._src.lax.control_flow.solves
- 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.special
- jax._src.lax.windowed_reductions
- jax._src.linear_util
- jax._src.maps
- jax._src.mesh
- jax._src.nn.functions
- jax._src.numpy.fft
- jax._src.numpy.lax_numpy
- jax._src.numpy.linalg
- jax._src.numpy.polynomial
- jax._src.numpy.reductions
- jax._src.numpy.setops
- jax._src.numpy.ufunc_api
- jax._src.numpy.ufuncs
- jax._src.numpy.vectorize
- jax._src.ops.scatter
- jax._src.ops.special
- jax._src.partition_spec
- jax._src.pjit
- jax._src.prng
- jax._src.profiler
- jax._src.random
- jax._src.scipy.fft
- jax._src.scipy.integrate
- 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.spatial.transform
- jax._src.scipy.special
- jax._src.scipy.stats._core
- jax._src.scipy.stats.bernoulli
- jax._src.scipy.stats.beta
- jax._src.scipy.stats.betabinom
- jax._src.scipy.stats.binom
- 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.gennorm
- jax._src.scipy.stats.geom
- jax._src.scipy.stats.kde
- jax._src.scipy.stats.laplace
- jax._src.scipy.stats.logistic
- jax._src.scipy.stats.multinomial
- jax._src.scipy.stats.multivariate_normal
- jax._src.scipy.stats.nbinom
- jax._src.scipy.stats.norm
- jax._src.scipy.stats.pareto
- jax._src.scipy.stats.poisson
- jax._src.scipy.stats.t
- jax._src.scipy.stats.truncnorm
- jax._src.scipy.stats.uniform
- jax._src.scipy.stats.vonmises
- jax._src.scipy.stats.wrapcauchy
- jax._src.sharding
- jax._src.sharding_impls
- jax._src.stages
- jax._src.third_party.numpy.linalg
- jax._src.third_party.scipy.betaln
- jax._src.third_party.scipy.linalg
- jax._src.tree_util
- jax._src.xla_bridge
- jax.example_libraries.optimizers
- jax.example_libraries.stax
- jax.experimental.custom_partitioning
- jax.experimental.host_callback
- jax.experimental.jet
- jax.experimental.mesh_utils
- jax.experimental.multihost_utils
- jax.experimental.sparse._base
- jax.experimental.sparse.ad
- jax.experimental.sparse.api
- jax.experimental.sparse.bcoo
- jax.experimental.sparse.bcsr
- jax.experimental.sparse.coo
- jax.experimental.sparse.csr
- jax.experimental.sparse.linalg
- jax.experimental.sparse.random
- jax.experimental.sparse.transform
- jax.experimental.x64_context
- jaxlib.xla_extension
- ml_dtypes._finfo
- ml_dtypes._iinfo
- numpy.exceptions