扩展指南

用于扩展 JAX 功能以及构建使用或与 JAX 接口的库的指南。

可扩展的 JAX 内部原理

• 在 JAX 中编写自定义 Jaxpr 解释器
• 使用 C++ 和 CUDA 的 GPU 自定义操作
• 检查正确性
• jax.extend 模块

库和扩展

• 基于 JAX 构建