扩展指南

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

可扩展的 JAX 内部机制

• JAX原语是如何工作的
• 理解 Jaxprs
• 在JAX中编写自定义Jaxpr解释器
• 使用 C++ 和 CUDA 进行 GPU 的自定义操作
• 检查正确性
• jax.extend 模块

库和扩展

• 基于 JAX 构建