Pallas: 一个 JAX 内核语言#
Pallas 是 JAX 的一个扩展，它使得为 GPU 和 TPU 编写自定义内核成为可能。它的目标是提供对生成代码的细粒度控制，同时结合 JAX 追踪和 jax.numpy API 的高级人体工程学。
本节包含使用 Pallas 的教程、指南和示例。另请参阅 jax.experimental.pallas 模块的 API 文档。
警告
Pallas 是实验性的，并且经常在变化。请参阅 Pallas 更新日志 了解最近的更改。
你可能会遇到错误和未实现的情况，例如，在降低需要仿真的高级 JAX 概念时，或者仅仅因为 Pallas 仍在开发中。
指南