Pallas: JAX 内核语言#

Pallas 是 JAX 的一个扩展，可以为 GPU 和 TPU 编写自定义内核。它的目标是对生成的代码提供细粒度的控制，并结合 JAX 跟踪和 jax.numpy API 的高级人体工程学。

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

警告

Pallas 是实验性的，并且经常更改。有关最近的更改，请参阅 Pallas 变更日志。

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

指南

TPU 后端指南

Mosaic GPU 后端指南

设计注释

其他