扩展指南

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

可扩展的 JAX 内部机制

• 在 JAX 中编写自定义 Jaxpr 解释器
• jax.extend 模块

库与扩展

• 基于 JAX 构建