Skip to main content
Ctrl+K
JAX documentation - Home

Getting started

  • Installation
  • Quickstart
  • Tutorials
    • Quickstart
    • Key concepts
    • Just-in-time compilation
    • Automatic vectorization
    • Automatic differentiation
    • Introduction to debugging
    • Pseudorandom numbers
    • Working with pytrees
    • Introduction to parallel programming
    • Stateful computations
    • Control flow and logical operators with JIT
    • Advanced automatic differentiation
    • External callbacks
    • Gradient checkpointing with jax.checkpoint (jax.remat)
    • JAX Internals: primitives
    • JAX internals: The jaxpr language
  • 🔪 JAX - The Sharp Bits 🔪
  • Frequently asked questions (FAQ)

More guides/resources

  • User guides
    • How to think in JAX
    • Profiling computation
    • Profiling device memory
    • Debugging runtime values
      • Compiled prints and breakpoints
      • The checkify transformation
      • JAX debugging flags
    • GPU performance tips
    • Persistent compilation cache
    • Pytrees
    • Errors
    • Ahead-of-time lowering and compilation
    • Exporting and serialization
      • Exporting and serializing staged-out computations
      • Shape polymorphism
      • Interoperation with TensorFlow
    • Type promotion semantics
    • Transfer guard
    • Pallas: a JAX kernel language
      • Pallas Quickstart
      • Grids and BlockSpecs
      • Pallas TPU
        • Writing TPU kernels with Pallas
        • Pipelining
        • Matrix Multiplication
        • Scalar Prefetch and Block-Sparse Computation
        • Distributed Computing in Pallas for TPUs
      • Pallas Design Notes
        • Pallas Design
        • Pallas Async Operations
      • Pallas Changelog
    • Foreign function interface (FFI)
    • Training a simple neural network, with tensorflow/datasets data loading
    • Training a simple neural network, with PyTorch data loading
    • Autobatching for Bayesian inference
  • Advanced guides
    • Distributed arrays and automatic parallelization
    • Manual parallelism with shard_map
    • Multi-host and multi-process environments
    • Distributed data loading
    • The Autodiff Cookbook
    • Custom derivative rules
    • Control autodiff’s saved values with jax.checkpoint (aka jax.remat)
    • Generalized convolutions in JAX
    • List of XLA compiler flags
  • Developer notes
    • Contributing to JAX
    • Building from source
    • Investigating a regression
    • Autodidax: JAX core from scratch
    • JAX Enhancement Proposals (JEPs)
      • 263: JAX PRNG Design
      • 2026: Custom JVP/VJP rules for JAX-transformable functions
      • 4008: Custom VJP and `nondiff_argnums` update
      • 4410: Omnistaging
      • 9263: Typed keys & pluggable RNGs
      • 9407: Design of Type Promotion Semantics for JAX
      • 9419: Jax and Jaxlib versioning
      • 10657: Sequencing side-effects in JAX
      • 11830: `jax.remat` / `jax.checkpoint` new implementation
      • 12049: Type Annotation Roadmap for JAX
      • 14273: `shard_map` (`shmap`) for simple per-device code
      • 15856: `jax.extend`, an extensions module
      • 17111: Efficient transposition of `shard_map` (and other maps)
      • 18137: Scope of JAX NumPy & SciPy Wrappers
      • 25516: Effort-based versioning
  • Extension guides
    • Writing custom Jaxpr interpreters in JAX
    • Custom operations for GPUs with C++ and CUDA
    • jax.extend module
      • jax.extend.core module
      • jax.extend.linear_util module
      • jax.extend.mlir module
      • jax.extend.random module
    • Building on JAX
  • Notes
    • API compatibility
    • Python and NumPy version support policy
    • jax.Array migration
    • Asynchronous dispatch
    • Concurrency
    • GPU memory allocation
    • Rank promotion warning
  • Public API: jax package
    • jax.numpy module
      • jax.numpy.fft.fft
      • jax.numpy.fft.fft2
      • jax.numpy.fft.fftfreq
      • jax.numpy.fft.fftn
      • jax.numpy.fft.fftshift
      • jax.numpy.fft.hfft
      • jax.numpy.fft.ifft
      • jax.numpy.fft.ifft2
      • jax.numpy.fft.ifftn
      • jax.numpy.fft.ifftshift
      • jax.numpy.fft.ihfft
      • jax.numpy.fft.irfft
      • jax.numpy.fft.irfft2
      • jax.numpy.fft.irfftn
      • jax.numpy.fft.rfft
      • jax.numpy.fft.rfft2
      • jax.numpy.fft.rfftfreq
      • jax.numpy.fft.rfftn
    • jax.scipy module
      • jax.scipy.stats.bernoulli.logpmf
      • jax.scipy.stats.bernoulli.pmf
      • jax.scipy.stats.bernoulli.cdf
      • jax.scipy.stats.bernoulli.ppf
    • jax.lax module
    • jax.random module
    • jax.sharding module
    • jax.debug module
    • jax.dlpack module
    • jax.distributed module
    • jax.dtypes module
    • jax.ffi module
    • jax.flatten_util module
    • jax.image module
    • jax.nn module
      • jax.nn.initializers module
    • jax.ops module
    • jax.profiler module
    • jax.stages module
    • jax.tree module
    • jax.tree_util module
    • jax.typing module
    • jax.export module
    • jax.extend module
      • jax.extend.core module
      • jax.extend.linear_util module
      • jax.extend.mlir module
      • jax.extend.random module
    • jax.example_libraries module
      • jax.example_libraries.optimizers module
      • jax.example_libraries.stax module
    • jax.experimental module
      • jax.experimental.checkify module
      • jax.experimental.compilation_cache module
      • jax.experimental.custom_partitioning module
      • jax.experimental.jet module
      • jax.experimental.key_reuse module
      • jax.experimental.mesh_utils module
      • jax.experimental.multihost_utils module
      • jax.experimental.pallas module
        • jax.experimental.pallas.mosaic_gpu module
        • jax.experimental.pallas.triton module
        • jax.experimental.pallas.tpu module
      • jax.experimental.pjit module
      • jax.experimental.serialize_executable module
      • jax.experimental.shard_map module
      • jax.experimental.sparse module
        • jax.experimental.sparse.BCOO
        • jax.experimental.sparse.bcoo_broadcast_in_dim
        • jax.experimental.sparse.bcoo_concatenate
        • jax.experimental.sparse.bcoo_dot_general
        • jax.experimental.sparse.bcoo_dot_general_sampled
        • jax.experimental.sparse.bcoo_dynamic_slice
        • jax.experimental.sparse.bcoo_extract
        • jax.experimental.sparse.bcoo_fromdense
        • jax.experimental.sparse.bcoo_gather
        • jax.experimental.sparse.bcoo_multiply_dense
        • jax.experimental.sparse.bcoo_multiply_sparse
        • jax.experimental.sparse.bcoo_update_layout
        • jax.experimental.sparse.bcoo_reduce_sum
        • jax.experimental.sparse.bcoo_reshape
        • jax.experimental.sparse.bcoo_slice
        • jax.experimental.sparse.bcoo_sort_indices
        • jax.experimental.sparse.bcoo_squeeze
        • jax.experimental.sparse.bcoo_sum_duplicates
        • jax.experimental.sparse.bcoo_todense
        • jax.experimental.sparse.bcoo_transpose
    • jax.lib module
    • jax.Array.addressable_shards
    • jax.Array.all
    • jax.Array.any
    • jax.Array.argmax
    • jax.Array.argmin
    • jax.Array.argpartition
    • jax.Array.argsort
    • jax.Array.astype
    • jax.Array.at
    • jax.Array.choose
    • jax.Array.clip
    • jax.Array.compress
    • jax.Array.committed
    • jax.Array.conj
    • jax.Array.conjugate
    • jax.Array.copy
    • jax.Array.copy_to_host_async
    • jax.Array.cumprod
    • jax.Array.cumsum
    • jax.Array.device
    • jax.Array.diagonal
    • jax.Array.dot
    • jax.Array.dtype
    • jax.Array.flat
    • jax.Array.flatten
    • jax.Array.global_shards
    • jax.Array.imag
    • jax.Array.is_fully_addressable
    • jax.Array.is_fully_replicated
    • jax.Array.item
    • jax.Array.itemsize
    • jax.Array.max
    • jax.Array.mean
    • jax.Array.min
    • jax.Array.nbytes
    • jax.Array.ndim
    • jax.Array.nonzero
    • jax.Array.prod
    • jax.Array.ptp
    • jax.Array.ravel
    • jax.Array.real
    • jax.Array.repeat
    • jax.Array.reshape
    • jax.Array.round
    • jax.Array.searchsorted
    • jax.Array.shape
    • jax.Array.sharding
    • jax.Array.size
    • jax.Array.sort
    • jax.Array.squeeze
    • jax.Array.std
    • jax.Array.sum
    • jax.Array.swapaxes
    • jax.Array.take
    • jax.Array.to_device
    • jax.Array.trace
    • jax.Array.transpose
    • jax.Array.var
    • jax.Array.view
    • jax.Array.T
    • jax.Array.mT
  • About the project
  • Change log
  • Glossary of terms
  • .rst

Pallas Design Notes

Pallas Design Notes#

Design

  • Pallas Design
    • Introduction
    • Pallas: Extending JAX for kernels
  • Pallas Async Operations
    • Background + Motivation
    • How do we implement decomposed Pallas async operations (on TPU)?
    • Why doesn’t this work?
    • Revenge of the State
    • Conclusion

previous

Distributed Computing in Pallas for TPUs

next

Pallas Design

By The JAX authors

© Copyright 2024, The JAX Authors.