Pallas: a JAX kernel language

Pallas is an extension to JAX that enables writing custom kernels for GPU and TPU. This section contains tutorials, guides and examples for using Pallas.

Guides

• Pallas Design
  • Introduction
  • Pallas: Extending JAX for kernels
• Pallas Quickstart
  • Hello world in Pallas
  • Pallas programming model

Platform Features

• Pallas TPU
  • Writing TPU kernels with Pallas
  • Pipelining and BlockSpecs