Pallas: a JAX kernel language

Pallas is an extension to JAX that enables writing custom kernels for GPU and TPU. It aims to provide fine-grained control over the generated code, combined with the high-level ergonomics of JAX tracing and the jax.numpy API.

This section contains tutorials, guides and examples for using Pallas. See also the jax.experimental.pallas module API documentation.

Warning

Pallas is experimental and is changing frequently. See the Pallas Changelog for the recent changes.

You can expect to encounter errors and unimplemented cases, e.g., when lowering of high-level JAX concepts that would require emulation, or simply because Pallas is still under development.

Guides

• Pallas Quickstart
  • Hello world in Pallas
  • Pallas programming model
• Grids and BlockSpecs
  • grid, a.k.a. kernels in a loop
  • BlockSpec, a.k.a. how to chunk up inputs

Platform Features

• Pallas TPU
  • Writing TPU kernels with Pallas
  • Pipelining
  • Matrix Multiplication
  • Scalar Prefetch and Block-Sparse Computation
  • Distributed Computing in Pallas for TPUs

Design Notes

• Pallas Design Notes
  • Pallas Design
  • Pallas Async Operations

Other

• Pallas Changelog