用 hijax 为 JAX 定义全新类型

Defining new Jax types with hijax

JAX 默认以数组为核心，但处理复杂聚合数据时，Pytree 的透明化往往不够用。本文介绍如何用 hijax 定义全新的 JAX 类型，让数据在 jaxpr 中保持单一身份，拥有独立的不变量、切线类型和批处理规则。以量化数组为例，展示了如何绕过 Pytree 限制，实现更精准的自动微分和显式分片。如果你想在 JAX 中构建更高级的数据结构，hijax 类型是实验性但强大的选择。

"有时透明化恰恰是你不想要的，某些数据最好被建模为拥有独立身份的新类型。"