用hijax为Jax定义新类型

背景

JAX 的核心抽象是数组：所有变换函数都接收数组、输出数组，追踪器看到的每个中间值都带有数组类型（如 f32[3,4]）。当需要处理聚合数据时，常规工具是 pytree：将数组打包成容器，JAX 在每次边界处自动将容器展平为数组叶子节点。然而，这种透明性有时恰恰是你不想要的。某些数据更适合建模为具有独立身份的新类型：

• 它在 jaxpr 中应表现为单一类型的单值，而不是一堆数组叶子；
• 它包含内部不变量，用户只能通过固定的操作集合创建和消费它，而不能自由构造或模式匹配其组件；
• 它的切向类型可能与原始结构不同——导数不应只是“同样的 pytree，但对应切向”；
• 在 vmap 下它可能有自己的批处理概念；
• 它可以在类型中携带分片信息，参与 JAX 的显式分片模式。

Hijax 类型（或“hi types”）提供了这种能力。用户通过继承 HiType 定义类型，注册一个 Python 类作为该类型的值载体，并编写输入/输出类型提及新类型的 hijax 原语。本文以量化数组类型为例，完整展示整个流程。

核心内容

TL;DR 速览

1. 继承 HiType 并实现 lo_ty、lower_val、raise_val，说明该类型及其值如何降级为普通（“lojax”）数组，然后调用 register_hitype 将值类与类型关联。
2. 编写 VJPHiPrimitive 子类，其 in_avals/out_aval 提及新类型；这些原语是产生和消费该类型值的唯一途径。
3. 自动微分：在类型上实现 to_tangent_aval，在原语上实现 VJP/JVP 规则。
4. vmap：在类型上实现 dec_rank 和 inc_rank，并设计自定义的 MappingSpec 子类，在原语上实现 batch 规则。映射过的 hi 类型参数需要显式指定 axis_size 和 in_axes/out_axes 项（值为 spec 类型）。
5. 类型内分片（显式模式）：在类型上记录分片数据（如 NamedSharding 字段），在 lo_ty 中消费它，并在原语的类型规则中传播它。

示例：量化数组

假设我们想使用量化到 int8 的数组。一个量化数组实际上是一对数组：int8 值和一个浮点尺度（scale），每行共享一个尺度（即沿最后一轴量化，每行一个尺度，如常见的按行/按通道量化方案）：

import os
os.environ["XLA_FLAGS"] = '--xla_force_host_platform_device_count=8'
from dataclasses import dataclass
import jax
import jax.numpy as jnp

@dataclass(frozen=True)
class QArray:
    qvalue: jax.Array  # int8[*leading, n]
    scale: jax.Array   # f32[*leading]

我们可以将 QArray 注册为 pytree 并完成工作。但思考一下我们会放弃什么：

• 不变量：两个组件是耦合的——scale 的形状必须等于 qvalue 去掉最后一轴，且 qvalue 只有与 scale 一起才有意义。作为 pytree，没有什么能阻止代码破坏这种关联；在变换下，JAX 本身只看到独立的叶子节点。
• jaxpr 中的类型：作为 pytree，量化数组在追踪代码中表现为两个不相关的数组值。我们更希望看到单一类型的单一值，使 jaxpr 能表达其含义。
• 切向：量化数组的值位于离散网格上，因此沿着网格扰动它们没有意义。但 pytree 的切向类型被迫等于其叶子切向类型的 pytree——而像 qvalue 这样的整数数组的切向类型是 float0 数组，只能携带微不足道的载荷。因此，作为 pytree，量化数组完全无法接受有用的扰动。我们想要的是为整个量化数组选择一个切向类型，比如其量化值所近似的连续 f32 数组。

因此，我们将 QArray 制作成 hijax 类型。

类型定义

Hijax 类型是 HiType 的子类。核心必需部分很小：

• lo_ty：说明该类型由哪些 lojax（数组）类型组成；
• lower_val 和 raise_val：将值转换为数组列表，以及从数组列表还原；
• 类型必须是可哈希且可比较相等的（frozen dataclass 同时满足这两点）。

这类似于 pytree 展平/反展平接口，但位于类型层面：仅给定类型，JAX 就能计算出降级后的类型，而无需值在手。

我们还在类型中添加一个 sharding 字段，记录值如何跨设备分区；在本文末尾的分片章节之前它可以被忽略。我们复用 JAX 的 NamedSharding 来描述 qvalue 组件的分区，并通过去掉最后一轴推导出 scale 的分区。该字段对 JAX 没有特殊意义——JAX 从不解释它：只有我们自己的方法才会消费它，主要是 lo_ty，它给组件类型打上分片标签。你也可以用自己设计的对象来跟踪类型上的分片信息，只要以相同方式消费即可。

from jax.experimental.hijax import HiType, ShapedArray, register_hitype
from jax.sharding import NamedSharding

@dataclass(frozen=True)
class QArrayTy(HiType):
    shape: tuple[int, ...]
    sharding: NamedSharding  # qvalue 的分片；scale 的分片从它推导

    # 降级：该类型由哪些数组类型组成，以及如何转换值
    def lo_ty(self):
        scale_sharding = self.sharding.update(spec=jax.P(*self.sharding.spec[:-1]))
        return [ShapedArray(self.shape, jnp.dtype('int8'), sharding=self.sharding),
                ShapedArray(self.shape[:-1], jnp.dtype('float32'), sharding=scale_sharding)]

    def lower_val(self, q):
        return [q.qvalue