Pallas：一个 JAX 内核语言#

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

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

警告

Pallas 尚处于实验阶段，并且经常发生变化。请参阅 Pallas 更新日志 了解最近的更改。

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

指南

平台特性

设计说明

其他