Getting Started
Further Resources
jax.debug.print
jax.debug.breakpoint
checkify
BlockSpec
shard_map
xmap
jax.checkpoint
jax.remat
jax.numpy
jax.scipy
jax.lax
jax.random
jax.sharding
jax.debug
jax.dlpack
jax.distributed
jax.dtypes
jax.flatten_util
jax.image
jax.nn
jax.nn.initializers
jax.ops
jax.profiler
jax.stages
jax.tree
jax.tree_util
jax.typing
jax.extend
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.array_api
jax.experimental.checkify
jax.experimental.host_callback
jax.experimental.maps
jax.experimental.pjit
jax.experimental.sparse
jax.experimental.jet
jax.experimental.custom_partitioning
jax.experimental.multihost_utils
jax.experimental.compilation_cache
jax.experimental.key_reuse
jax.experimental.mesh_utils
jax.lib
Jaxpr(constvars, invars, outvars, eqns[, ...])
Jaxpr
ClosedJaxpr(jaxpr, consts)
ClosedJaxpr
