擴充指南

擴充 JAX 功能以及建構使用或介接 JAX 的函式庫的指南。

可擴充的 JAX 內部機制

• 在 JAX 中撰寫自訂 Jaxpr 解譯器
• 使用 C++ 和 CUDA 為 GPU 撰寫自訂運算
• 檢查正確性
• jax.extend 模組

函式庫和擴充功能

• 基於 JAX 建構