Pallas：一种 JAX 内核语言

Pallas 是 JAX 的一个扩展，允许为 GPU 和 TPU 编写自定义内核。它旨在提供对生成的代码的精细控制，同时结合 JAX 跟踪和 jax.numpy API 的高级人体工程学。

本节包含有关使用 Pallas 的教程、指南和示例。另请参阅 jax.experimental.pallas 模块 API 文档。

警告

Pallas 正在进行实验并且经常变化。请参阅 Pallas 更新日志 以了解最近的更改。

您可能会遇到错误和未实现的情况，例如，当降低需要模拟的高级 JAX 概念时，或者仅仅因为 Pallas 仍在开发中。

指南

• Pallas 快速入门
  • Pallas 中的你好世界
  • Pallas 编程模型
• 网格和 BlockSpecs
  • grid，又名循环中的内核
  • BlockSpec，又名如何分割输入

平台特性

• Pallas TPU
  • 使用 Pallas 编写 TPU 内核
  • 流水线
  • 矩阵乘法
  • 标量预取和块稀疏计算
  • Pallas 中用于 TPU 的分布式计算

设计说明

• Pallas 设计说明
  • Pallas 设计
  • Pallas 异步操作

其他

• Pallas 更新日志