Pallas TPU

TPU specific documentation.

Guides

• Writing TPU kernels with Pallas
  • What is a TPU?
  • Noteworthy properties and restrictions
  • Supported operations
• Pipelining and BlockSpecs
  • TPU and its memory spaces
  • Constraints of using VMEM/SMEM
  • Primer: Pipelining
  • Pipelining in Pallas
  • Handling reductions
  • TPUs in Megacore configuration
  • Conclusion