扩展指南

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

可扩展的 JAX 内部组件

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

库和扩展

• 基于 JAX 构建