Skip to main content
Ctrl+K

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
  • Pallas Quickstart
  • Writing TPU kernels with Pallas

previous

Transfer guard

next

Pallas Design