Getting started
Resources, guides, and references
shard_map
jax.checkpoint
jax.remat
checkify
set_xla_metadata
Ref
jax.numpy
jax.scipy
jax.lax
jax.random
jax.sharding
jax.debug
jax.dlpack
jax.distributed
jax.dtypes
jax.ffi
jax.flatten_util
jax.image
jax.nn
jax.nn.initializers
jax.ops
jax.profiler
jax.ref
jax.stages
jax.test_util
jax.tree
jax.tree_util
jax.typing
jax.export
jax.extend
jax.extend.core
jax.extend.linear_util
jax.extend.mlir
jax.extend.random
jax.example_libraries
jax.example_libraries.optimizers
jax.example_libraries.stax
jax.experimental
jax.experimental.checkify
jax.experimental.compilation_cache
jax.experimental.custom_dce
jax.experimental.custom_partitioning
jax.experimental.jet
jax.experimental.key_reuse
jax.experimental.mesh_utils
jax.experimental.multihost_utils
jax.experimental.pallas
jax.experimental.pallas.mosaic_gpu
jax.experimental.pallas.triton
jax.experimental.pallas.tpu
jax.experimental.serialize_executable
jax.experimental.sparse
float0
DType class corresponding to the scalar type and dtype of the same name.
Please seenumpy.dtype for the typical way to createdtype instances andData type objects (dtype) for additionalinformation.
