扩展指南# 扩展 JAX 功能以及构建使用或与 JAX 交互的库的指南。 可扩展的 JAX 内部 JAX 原语如何工作 理解 Jaxpr 在 JAX 中编写自定义 Jaxpr 解释器 使用 C++ 和 CUDA 为 GPU 创建自定义操作 检查正确性 jax.extend 模块 库和扩展 基于 JAX 构建
