使用 hijax 定义新的 JAX 类型 # JAX 的内置数据类型是数组：你转换的函数接受数组并输出数组，每个中间的追踪机制看到的都有数组类型，比如 f32[3,4]。当你希望处理聚合数据时，通常使用的工具是 pytree：你将数组打包到容器中，JAX 在每个边界透明地将包裹压缩为其数组叶子。但有时你确实不想要透明性。一些数据最好建模为一种新类型，具有其自身的身份：它在 jaxprs 中应该作为单一类型的单一值出现，而不是作为一串数组叶子；它有内部的不变性，因此用户应该只通过固定的操作集生成和消费它，而不是自由地构造或模式匹配其组成部分；它的切线类型可能与其原始结构不同，因此关于它的导数不仅仅是“同样的 pytree，但针对切线”；它可能在 vmap 下有自己的批处理概念；它可以在类型中承载分片信息，参与 JAX 的显式分片模式。Hijax 类型（或称“hi 类型”）提供了这种功能。你通过子类化 HiType 来定义类型，注册一个 Python 类来承载该类型的值，并编写输入和输出类型提到新类型的 hijax 原语。本文件通过一个运行示例走完整个过程：一个量化数组类型。我们假设对 hijax 原语有一定的熟悉度；有关它们的介绍，请参见 使用 hijax 原语的自定义导数规则。像 hijax 的其他所有内容一样，这都是实验性的：期待从 jax.experimental.hijax 导入，并期待 API 逐步演变。 简而言之 # 子类化 HiType 并实现 lo_ty、lower_val 和 raise_val 来说明类型及其值如何降低为普通（“lojax”）数组，然后调用 register_hitype 来将你的值类与类型关联。编写 VJPHiPrimitive 子类，其 in_avals / out_aval 提到新类型；这些是生成和消费该类型值的唯一方式。对于自动微分，实现 to_tangent_aval 在类型上，以及 VJP/JVP 规则在原语上。对于 vmap，分别在类型上实现 dec_rank 和 inc_rank，及你自己设计的 MappingSpec 子类，并在原语上实现批处理规则。被映射的 hi 类型参数需要显式的 axis_size 和 spec-valued in_axes / out_axes 条目。对于类型中的分片（显式模式），在你的类型上记录分片数据（例如 NamedSharding 字段），在 lo_ty 中消费它，并在你的原语的类型规则中传播。 示例：量化数组 # 假设我们想处理量化为 int8 的数组。量化数组实际上是一对数组：int8 值和每行共享的浮点缩放（即，我们沿最后一个轴量化，每行一个缩放，如常见的逐行/逐通道量化方案）： import os os.environ["XLA_FLAGS"] = '--xla_force_host_platform_device_count=8' # （8 个 CPU 设备，用于最后的分片部分） 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 本身只看到独立的叶子。类型在 jaxprs 中。作为 pytree，量化数组在追踪代码中表现为两个无关的数组值。我们更希望看到一个值，一个类型，因此 jaxprs 说明它们的含义。切线。量化数组的值生活在一个离散的网格上，因此在网格上扰动它们毫无意义。但 pytree 的切线类型被迫是其叶子切线类型的 pytree——而像 qvalue 这样的整数数组的切线类型是一个 float0 数组，只能承载一个微不足道的有效载荷。因此，作为一个 pytree，量化数组根本无法接受有用的扰动。我们想为量化数组整体选择切线类型，例如量化值所近似的连续 f32 数组。因此我们将使 QArray 成为 hijax 类型。类型 # 一个 hijax 类型是 HiType 的子类。所需的核心很小：lo_ty 指明哪些 lojax（数组）类型构成该类型；lower_val 和 raise_val 分别将值转换为和从该数组列表中转换；该类型必须是可哈希和可比较的（一个冻结的数据类同时提供这两者）。这像 pytree 的压平/反压平接口，但它存在于类型的级别：仅给定类型，JAX 可以计算降低类型，而无需手头有值。我们还给类型一个分片字段，记录值如何在设备间分区；在文档末尾的分片部分之前，它可以被忽略。我们重用 JAX 的 NamedSharding 来描述 t 的分区。