Pallas：JAX 核心語言#

Pallas 是 JAX 的擴充功能，可讓您為 GPU 和 TPU 撰寫自訂核心。其目標是提供對產生程式碼的細緻控制，並結合 JAX 追蹤和 jax.numpy API 的高階人體工學。

本節包含使用 Pallas 的教學、指南和範例。另請參閱 jax.experimental.pallas 模組 API 文件。

警告

Pallas 屬於實驗性質，且變更頻繁。請參閱 Pallas 變更日誌 以了解近期變更。

您可以預期會遇到錯誤和未實作的情況，例如，當降低需要模擬的高階 JAX 概念時，或僅僅因為 Pallas 仍在開發中。

指南

平台功能

設計筆記

其他