Hijax を用いて新しい JAX の型を定義する

はじめに

JAX の内部で最も基本的なデータ構造は**「アレイ（配列）」**です。変換する関数はアレイを受け取り、アレイを生み出します。追跡機構がすべての中間段階で見つけるものも、すべて

f32[3,4]

集合データを取り扱う場合、通常は「ピートリー（pytree）」というツールが使われます。これとは逆に、ピートリーでは配列をコンテナに束ねていますが、JAX は境界で透明にその束ね物をアレイの葉（leaves）としてフラット化します。

しかし、透明さが望ましくないこともあります。一部のデータは独自のアイデンティティを持つ新しい型としてモデル化する方が適しています。具体的には以下のような特徴を持つ必要があります：

• 追跡されたコード（

  jaxprs

  ） では、アレイの葉という形で散在しているのではなく、単一の型の単一値として表示されるべきです。
• 内部の不変条件（invariants）を有するため、ユーザーはその成分を自由に構築したりパターンマッチングを行ったりするのではなく、固定された一連の操作を通じてのみ生成・消費するべきです。
• その接線型（tangent type）が元の構造とは異なる場合があります。これにより、それに対する微分が単に「タンジェントを持つ同じピートリー」では済まず、より一般的なケースを扱えるようになります。

• vmap

  下で独自のバッチングの概念を有する可能性があります。
• JAX の明示的なシェーディング（sharding）モードにおいて、型の中にシェーディング情報を保持でき、参与其中できます。

Hijax 型（または「hi types」）がこの機能を提供します。

HiType

Hijax プライミティブに対するいくつかの基本的な知識があると仮定しますが、それらの概要については「Hijax プライミティブを用いたカスタム微分則」を参照してください。Hijax のすべての機能と同様に、これは実験的であり、

jax.experimental.hijax

TL;DR

• HiType

  をサブクラス化し、型とその値が通常の（"lojax"）アレイにどのように低下（lower）するかを記述する

  lo_ty

  ,

  lower_val

  , および

  raise_val

  を実装し、次に

  register_hitype

  を呼び出して、あなたの値クラスをあなたの型に関連付けます。
• 新しい型が言及された入出力型を持つ VJPHiPrimitive サブクラスのプライミティブを記述します。これが型の値が生産され消費される唯一の方法です。
• 自動微分のためには、接線型を指定する

  to_tangent_aval

  を型に実装し、プライミティブ上では VJP/JVP 則を実装します。
• vmap のためは、型のランクの変更を行ったり（

  dec_rank

  および

  inc_rank

  ）、独自の設計を持つバッチ処理仕様（

  MappingSpec

  サブクラス）やバッチ処理の則をプライミティブに実装します。Hi 型アグリメントの上でのマッピングには、明示的な

  axis_size

  と仕様値付きの

  in_axes

  /

  out_axes

  エントリが必要です。
• 型のシェーディングについて（明示モード）、型上にシェーディングデータを記録（例：

  NamedSharding

  フィールド）し、

  lo_ty

  でそれを読み込み、プライミティブのタイピング則でそれを伝播させます。

例：定量化アレイ

int8 に量子化されたアレイを扱うことを考えましょう。定量化アレイは本質的に 2 つのアレイのペアです：int8 の値と、各行で共有される浮動小数点スケールです（つまり、最後の軸に沿って量子化し、行ごとに 1 つのスケールを持つ、一般的な行ごと/チャンネルごとの量子化スキームと同様）：

import os
os.environ["XLA_FLAGS"] = '--xla_force_host_platform_device_count=8'  # (最後に示すシェーディング節のために 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

1. 不変条件。 2 つの成分は連結しています：

   scale

   は

   qvalue

   の最後の軸を除いた形状を持たなければならず、

   qvalue

   はスケールと組み合わさって初めて意味を持ちます。ピートリーとしてしまうと、コードがこれらのストリームを横断することを妨げる何ものもなく、変換下では JAX 自体は単に独立した葉しか見えません。

2. jaxprs

   内の型。 ピートリーとして、定量化アレイは追跡されたコードでは無関係の 2 つのアレイ値として表示されます。私たちは

   qvalue

   ,

   scale

   のように関連付けられていることが示される、単一値で単一の型の表現を好みます。
3. 接線型。 定量化アレイの値は離散的なグリッド上に存在するため、そのグリッドに沿って摂動させることには意味がありません。しかし、ピートリーの接線型は、必ず葉の接線型のピートリーでなければならないという強制があり、

   qvalue

   といった整数アレイの接線型の型は浮動小数点ゼロ（float0）のアレイであり、これは非自明なペイロードを運ぶことができます。したがって、ピートリーとして定量化アレイは有用な摂動を一切認めません。私たちは、定量化アレイ全体に対して接線型を選択することを望み、例えばその量子化された値が近似する連続的な f32 アレイのようなものです。

そこで代わりに、

QArray

型の定義

Hijax 型は

HiType

• lo_ty

  は、この型を構成する lojax（アレイ）の型を示します；

• lower_val

  と

  raise_val

  は値をそのアレイリストに、およびその逆に変換します；
• 型はイテラビリティと等価比較が可能でなければなりません（フrozen データクラスはこれらをくれます）。

これはピートリーの

flatten

unflatten

また、型にはパーティション情報が記録されたシェーディングフィールドを与えます（最後に示すシェーディング節で無視可能）。ここでは JAX の

NamedSharding

qvalue

scale

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 のそれはそこから導出される

    # lowering: この型を構成するアレイの型と、値の変換方法
    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, q.scale]
    
    def raise_val(self, qvalue, scale):
        return QArray(qvalue, scale)

    # autodiff: 定量化アレイの接線型は単なる浮動小数点アレイ（後で参照）
    def to_tangent_aval(self):
        return ShapedArray(self.shape, jnp.dtype('float32'),
                           sharding=self.sharding)

    # 印刷、例：jaxprs 内
    def str_short(self, short_dtypes=False, mesh_axis_types=False):
        dims = [str(d) if p is None else f'{d}@{p}'
                for d, p in zip(self.shape, self.sharding.spec)]
        return f'q8[{",".join(dims)}]'
    __repr__ = str_short

register_hitype(QArray, lambda q: QArrayTy(q.qvalue.shape,
                                           jax.typeof(q.qvalue).sharding))

register_hitype

jax.typeof

ShapedArray

qvalue

jax.typeof

QArrays

プライミティブ

ピートリーの場合、ユーザーは自由に値を構築・分解できます。Hijax 型の場合、値は新しい型を言及する declared types を持つ Hijax プライミティブによってのみ生産され消費されます。ここで不変条件が強制されます：すべてのプライミティブがそれらを保存すれば、常に成り立つからです。

私たちの 2 つのプライミティブは

quantize

dequantize

expand

from jax.experimental.hijax import VJPHiPrimitive

class Quantize(VJPHiPrimitive):
    def __init__(self, x_aval):
        if x_aval.dtype != jnp.dtype('float32'): raise TypeError(x_aval.dtype)
        self.in_avals = (x_aval,)
        self.out_aval = QArrayTy(x_aval.shape, x_aval.sharding)
        self.params = {}
        super().__init__()

    def expand(self, x):
        scale = jnp.max(jnp.abs(x), axis=-1) / 127.
        qvalue = jnp.round(x / scale[..., None]).astype(jnp.int8)
        return QArray(qvalue, scale)

    # ストレートスルー推定器：恒等写像であるかのように微分
    def vjp_fwd(self, nzs_in, x):
        return self(x), None

    def vjp_bwd_retval(self, _res, g):
        return (g,)

class Dequantize(VJPHiPrimitive):
    def __init__(self, q_aval):
        self.in_avals = (q_aval,)
        self.out_aval = ShapedArray(q_aval.shape, jnp.dtype('float32'),
                                    sharding=q_aval.sharding)
        self.params = {}
        super().__init__()

    def expand(self, qx):
        return qx.qvalue.astype('float32') * qx.scale[..., None]

    def vjp_fwd(self, nzs_in, qx):
        return self(qx), None

    def vjp_bwd_retval(self, _res, g):
        return (g,)

def quantize(x):
    return Quantize(jax.typeof(x))(x)

def dequantize(qx):
    return Dequantize(jax.typeof(qx))(qx)

Quantize

out_aval

Dequantize

in_avals

QArrayTys

expand

QArray

すべて eager（熱心）に機能します：

x = jnp.array([[1., 2., 3.],
               [4., -5., 6.]])

qx = quantize(x)
print(qx)
print(jax.typeof(qx))
print(dequantize(qx))

QArray(qvalue=Array([[  42,   85,  127],
       [  85, -106,  127]], dtype=int8), scale=Array([0.02362205, 0.04724409], dtype=float32))
q8[2,3]
[[ 0.992126  2.007874  3.      ]
 [ 4.015748 -5.007874  6.      ]]

プライミティブは外側、コンテナは内側

上記の

expand

QArray

expand

lower_val

raise_val

その理由は、追跡されたコードが実際に見ることからです。追跡の下では、定量化アレイは

QArray

q8[...]

Tracer

try:
    jax.jit(lambda qx: qx.qvalue)(qx)
except AttributeError as e:
    print('AttributeError:', e)

AttributeError: DynamicJaxprTracer has no attribute qvalue

そしてさらに悪いことに、コンストラクタを呼び出してもすぐに失敗しない場合があります。それはトレーサーを容器内に密かに運び込み、JAX はそれを一つの不透明な具体値として扱い、ミスは後に混乱するエラーとして表面化します（ここでは定数ハンドラーの欠如；grad 下ではリークしたトレーサーのエラー）：

def bad_quantize(x):
    scale = jnp.max(jnp.abs(x), axis=-1) / 127.
    return QArray(jnp.round(x / scale[..., None]).astype('int8'), scale)

try:
    jax.jit(bad_quantize)(x)
except TypeError as e:
    print('TypeError:', e)

TypeError: No constant handler for type: <class 'jax._src.interpreters.partial_eval.DynamicJaxprTracer'>

expand

expand

QArray

QArray

（ドキュメント内の他の場所でトップレベルで行われる属性の読み取り（例：

qx.qvalue

より現実的な演算：密度 × 定量化 matmul

変換演算だけでは API は薄いです。実際には、その型を直接消費する演算で定量化アレイ型は存在意義を築きます。典型的な例は、活性化

x

class MatmulQ(VJPHiPrimitive):
    def __init__(self, x_aval, q_aval):
        if not (isinstance(q_aval, QArrayTy) and len(x_aval.shape) == 2 and
                len(q_aval.shape) == 2 and x_aval.shape[1] == q_aval.shape[0]):
            raise TypeError(f'bad matmul_q operand types: {x_aval} @ {q_aval}')
        self.in_avals = (x_aval, q_aval)
        x_spec, q_spec = x_aval.sharding.spec, q_aval.sharding.spec
        if x_spec[1] is not None or q_spec[0] is not None:
            raise TypeError('matmul_q requires unsharded contraction axes, got '
                            f'{x_aval} @ {q_aval}')
        out_sharding = x_aval.sharding.update(spec=jax.P(x_spec[0], q_spec[1]))
        self.out_aval = ShapedArray((x_aval.shape[0], q_aval.shape[1]),
                                    jnp.dtype('float32'), sharding=out_sharding)
        self.params = {}
        super().__init__()

    def expand(self, x, qw):
        # 行ごとのスケールを密度演算子に畳み込み、その後で int8 ペイロードに対して
        # 直接的な matmul を適用する。
        return (x * qw.scale) @ qw.qvalue.astype(jnp.float32)

    def vjp_fwd(self, nzs_in, x, qw):
        return self(x, qw), (x, qw)

    def vjp_bwd_retval(self, res, g):
        x, qw = res
        w = dequantize(qw)   # 則は追跡されたコード：プライミティブを使用する。
        # コタンジェントはプライマル（primals）と同じ場所に存在するため、プライマル演算子の
        # シャディングを使用して（おそらくシャーディングされた）収縮を区別します。
        return (jnp.matmul(g, w.T, out_sharding=jax.typeof(x).sharding),
                jnp.matmul(x.T, g, out_sharding=jax.typeof(w).sharding))

def matmul_q(x, qw):
    return MatmulQ(jax.typeof(x), jax.typeof(qw))(x, qw)

いくつかのポイントに注目してください。型シグネチャはアレイ型と hi 型を自由に混ぜています：

in_avals

(ShapedArray, QArrayTy)

expand

また、直前の節での規律が機能していることも注目してください：

expand

qw.scale

qw.qvalue

dequantize

（タイピング則も出力シェーディングを計算します——出力行は

x

qw

out_sharding

w = jnp.arange(12., dtype='float32').reshape(3, 4) / 12.
qw = quantize(w)

print(matmul_q(x, qw))
print(x @ dequantize(qw))  # リファレンス

[[2.6627297 3.1706038 3.6587927 4.166667 ]
 [2.3077428 2.7447505 3.1463253 3.5833333]]
[[2.6627297 3.1706038 3.6587927 4.166667 ]
 [2.3077426 2.7447505 3.1463253 3.5833333]]

JAXprs 内の Hijax 型

追跡する際、定量化アレイは単一の型

q8[2,3]

jax.jit(lambda x: dequantize(quantize(x))).trace(x).jaxpr

{ lambda ; a:f32[2,3]. let
    b:q8[2,3] = call_hi_primitive[_prim=Quantize[{}]] a
    c:f32[2,3] = call_hi_primitive[_prim=Dequantize[{}]] b
  in (c,) }

ピートリーアプローチと比較すると、同じ計算でも、ペアに結びついているという表示がない状態で、4 つのアレイ型中間体が示されることになります。Hi 型は低下時のみ消え去り、

expand

q8[...]

lo_ty

混合演算子種の演算子は、直接同じように読み取れます——1 つの方程式で

f32[2,3] @ q8[3,4] -> f32[2,4]

jax.jit(matmul_q).trace(x, qw).jaxpr

（注：実際の出力は

x

qw

{ lambda ; a:f32[2,3] b:q8[3,4]. let
    c:f32[2,4] = call_hi_primitive[_prim=MatmulQ[{}]] a b
  in (c,) }

print(jax.jit(lambda x: dequantize(quantize(x)))(x))   # QArray internal

qx2 = jax.jit(quantize)(x)                             # QArray result
print(jax.typeof(qx2))

print(jax.jit(dequantize)(qx2))                        # QArray argument

[[ 0.992126  2.007874  3.      ]
 [ 4.015748 -5.007874  6.      ]]
q8[2,3]
[[ 0.992126  2.007874  3.      ]
 [ 4.015748 -5.007874  6.      ]]

自動微分と接線型

ここで Hijax 型がその存在価値を示します。型上で、私たちは実装しました：

def to_tangent_aval(self):
    return ShapedArray(self.shape, jnp.dtype('float32'))

これは、「定量化アレイの接線型は単なる f32 アレイである」と述べています。これはいかなるピートリーでも表現できません：ピートリーの接線型は、常にその葉の接線型のピートリーであり、int8 葉

qvalue

プライミティブ上のストレートスルー VJP 則と合わせると、勾配は量子化を通じて恒等写像のように流れます：

def f(x):
    return jnp.sum(dequantize(quantize(x)))

そして、定量化アレイ入力に関する微分は、接線型が規定する通り、単なる浮動小数点アレイを生み出します：

def g(qx):
    return jnp.sum(dequantize(qx) ** 2)

print(jax.grad(g)(qx))
print(jax.typeof(jax.grad(g)(qx)))

[[  1.984252   4.015748   6.      ]
 [  8.031496 -10.015748  12.      ]]
float32[2,3]

混合演算子種の演算子についても同様です（量子化された matmul など）：損失に関する両方の演算子に対する微分を行うと、密度活性化に対して f32 グリデントを、量子化された重みに対して f32 グリデントを与えます：

def loss(x, qw):
    return jnp.sum(matmul_q(x, qw) ** 2)

grad_x, grad_qw = jax.grad(loss, argnums=(0, 1))(x, qw)
print(jax.typeof(grad_x), jax.typeof(grad_qw))

float32[2,3] float32[3,4]

接線型を f32 アレイにすることは選択であり、ここに実際の設計スペースがあります。代わりに

QArrayTy

QArrayTy

vspace_zero

vspace_add

vmap とマッピング仕様

定量化アレイをマッピングすることは何を意味するのでしょうか？アレイの場合、vmap の

in_axes

out_axes

in_axes

out_axes

axis_size

私たちの定量化アレイの場合、行ごとのスケールのおかげで、QArrays のバッチは単に大きな QArray です：新しい先頭軸に沿って型

q8[2,3]

qvalue

(n, 2, 3)

scale

(n, 2)

q8[n,2,3]

from jax.experimental.hijax import MappingSpec

@dataclass(frozen=True)
class QArraySpec(MappingSpec):
    pass  # QArrays は先頭軸に沿ってのみマッピングされる

（仕様があなたの型が必要とするほど豊かになる可能性があります。タプルのような hi 型は、各コンポーネントごとに 1 つの軸を持つ仕様を使用する可能性があります；このドキュメントの最後にあるタプルの例を参照してください。）

型の上では、

dec_rank

inc_rank

def qarray_dec_rank(self, size, spec):
    assert isinstance(spec, QArraySpec) and self.shape[0] == size
    return QArrayTy(self.shape[1:],
                    self.sharding.update(spec=jax.P(*self.sharding.spec[1:])))

def qarray_inc_rank(self, size, spec):
    assert isinstance(spec, QArraySpec)
    return QArrayTy((size, *self.shape),
                    self.sharding.update(spec=jax.P(None, *self.sharding.spec)))

QArrayTy.dec_rank = qarray_dec_rank
QArrayTy.inc_rank = qarray_inc_rank

（私たちは段階的にクラスのメソッドにアタッチしています；実際のコードではこれらは単に

QArrayTy

プライミティブの上では、バッチ則を実装します。バッチ則は、バッチ化された引数とマッピング仕様（未バッチ化の引数に対しては

None

def quantize_batch(self, axis_data, args, in_dims):
    x, = args
    d, = in_dims
    if d is None:
        return quantize(x), None
    x = jnp.moveaxis(x, d, 0)
    return quantize(x), QArraySpec()
Quantize.batch = quantize_batch

def dequantize_batch(self, axis_data, args, in_dims):
    qx, = args
    d, = in_dims
    if d is None:
        return dequantize(qx), None
    assert isinstance(d, QArraySpec)
    return dequantize(qx), 0
Dequantize.batch = dequantize_batch

行ごとの量子化は任意のランクで適用されるため、両方の則ともスタックされた値に未バッチ化された操作を適用するだけで済みます——これは、同じ型ファミリーの値を持つバッチを持つ型の象徴です。

これで vmap が機能します。定量化アレイ出力へのマッピングでは、軸サイズは通常のように配列引数から推測され、

out_axes

xs = jnp.arange(24., dtype='float32').reshape(4, 2, 3)

qxs = jax.vmap(quantize, out_axes=QArraySpec())(xs)
print(jax.typeof(qxs))
print(qxs.qvalue.shape, qxs.scale.shape)

q8[4,2,3]
(4, 2, 3) (4, 2)

定量化アレイ入力上のマッピングでは、

in_axes

axis_size

xs_roundtrip = jax.vmap(dequantize, in_axes=QArraySpec(), axis_size=4)(qxs)
print(jax.typeof(xs_roundtrip))

通常のすべての合成が機能します——jit の vmap、

print(jax.typeof(jax.vmap(jax.jit(dequantize), in_axes=QArraySpec(),
                          axis_size=4)(qxs)))

grad の vmap、など：

def norm_quantized(x):
    return jnp.sum(dequantize(quantize(x)) ** 2)

print(jax.vmap(jax.grad(norm_quantized))(xs).shape)

スキャンと先頭軸

jax.lax.scan

leading_axis_spec

dec_rank

inc_rank

def qarray_leading_axis_spec(self):
    return QArraySpec()

QArrayTy.leading_axis_spec = qarray_leading_axis_spec

前のセクションからの定量化アレイのスタックをスキャンする際、身体は各ステップで 1 つの

q8[2,3]

axis_size

length

def sum_dequantized(total, qx):
    return total + jnp.sum(dequantize(qx)), ()

total, () = jax.lax.scan(sum_dequantized, 0., qxs, length=4)
print(total)

Hi 値はスタックされた出力としてもループのカリーとしても機能します。ここではカリーは定量化アレイであり、各蓄積ステップ後に再量子化され、2 つ目の出力は各ステップで新しい QArray を

q8[4,2,3]

def accum_quantized(qtotal, x):
    return quantize(dequantize(qtotal) + x), quantize(2 * x)

qzero = quantize(jnp.zeros((2, 3), 'float32'))
qtotal, qys = jax.lax.scan(accum_quantized, qzero, xs)
print(jax.typeof(qtotal))
print(jax.typeof(qys))

型のシェーディング：明示モード

最後に、シェーディング。JAX の明示的なシェーディングモード（並列性ガイドを参照）では、シェーディングはアレイ型の一部です：

jax.typeof

QArrayTy

定量化アレイをデバイス間でパーティション化することは何を意味するのでしょうか？コンポーネントは一緒に移動します：行をシャードする場合、各デバイスはそれらの行の量子化値とそのスケールを持っています。通常のように設計の選択があり、私たちは

lo_ty

qvalue

scale

それが機能するところを見てみましょう。メッシュを作ります（これはドキュメントの最初のセルで要求された 8 クロウデバイスの使用）：

mesh = jax.make_mesh((4,), ('i',))
jax.set_mesh(mesh)

rows = jax.device_put(jnp.arange(24., dtype='float32').reshape(8, 3),
                      jax.P('i'))

qrows = quantize(rows)
print(jax.typeof(qrows))
print(qrows.qvalue.sharding.spec, qrows.scale.sharding.spec)
print(jax.typeof(dequantize(qrows)))

q8[8@i,3]
P('i', None) P('i',)
float32[8@i,3]

追跡中も同様で、

jaxprs

out_aval

jax.jit(lambda x: dequantize(quantize(x))).trace(rows).jaxpr

{ lambda ; a:f32[8@i,3]. let
    b:q8[8@i,3] = call_hi_primitive[_prim=Quantize[{}]] a
    c:f32[8@i,3] = call_hi_primitive[_prim=Dequantize[{}]] b
  in (c,) }

定量化された matmul もシェーディングを伝播します：タイピング則は出力行を

x

qw

w2 = jnp.arange(12., dtype='float32').reshape(3, 4) / 12.

print(jax.typeof(matmul_q(rows, quantize(w2))))                # rows sharded

qw2 = quantize(jax.device_put(w2, jax.P(None, 'i')))
print(jax.typeof(qw2))                                         # cols sharded
print(jax.typeof(matmul_q(jnp.ones((2, 3), 'float32'), qw2)))

float32[8@i,4]
q8[3,4@i]
float32[2,4@i]

そして、収縮チェックは trace time で型エラーとして発火します：

xk = jax.device_put(jnp.ones((2, 8), 'float32'), jax.P(None, 'i'))
qk = quantize(jax.device_put(jnp.ones((8, 3), 'float32'), jax.P('i')))

try:
    matmul_q(xk, qk)
except TypeError as e:
    print('TypeError:', e)

TypeError: matmul_q requires unsharded contraction axes, got float32[2,8@i] @ q8[8@i,3]

自動微分はこれらすべてと組成します。

MatmulQ

out_sharding

qw

def qloss(x, qw):
    return jnp.sum(matmul_q(x, qw) ** 2)

grad_rows, grad_qw2 = jax.grad(qloss, argnums=(0, 1))(rows, quantize(w2))
print(jax.typeof(grad_rows), jax.typeof(grad_qw2))

float32[8@i,3] float32[3,4]

注意点として：JAX は hi プライミティブが宣言した出力シェーディングを

expand

out_aval

expand

Shard_map とパーティション仕様

明示モードはプログラムの全体ビューを保ちながら値をパーティション化します。その対照は

shard_map

shard_map

in_specs

out_specs

shard_map

HiPspec

MappingSpec

in_specs

out_specs

to_lo

jax.P

from jax.experimental.hijax import HiPspec

@dataclass(frozen=True)
class QArrayP(HiPspec):
    spec: jax.P  # 先頭軸のパーティション化；最後の軸はそのまま

    def to_lo(self):
        return (self.spec, self.spec)  # qvalue と scale は一緒にシャードされる

（

scale

qvalue

qvalue

型の上では、

shard

unshard

def qarray_shard(self, mesh, manual_axes, check_vma, spec):
    qvalue_ty, _ = self.lo_ty()
    qspec, _ = spec.to_lo()
    shard_ty = qvalue_ty.shard(mesh, manual_axes, check_vma, qspec)
    return QArrayTy(shard_ty.shape, shard_ty.sharding)

def qarray_unshard(self, mesh, check_vma, spec):
    qvalue_ty, _ = self.lo_ty()
    qspec, _ = spec.to_lo()
    full_ty = qvalue_ty.unshard(mesh, check_vma, qspec)
    return QArrayTy(full_ty.shape, full_ty.sharding)

QArrayTy.shard = qarray_shard
QArrayTy.unshard = qarray_unshard

これで定量化アレイは

shard_map

@jax.jit
@jax.shard_map(in_specs=jax.P('i'), out_specs=QArrayP(jax.P('i')))
def quantize_shards(x):
    assert jax.typeof(x).shape == (2, 3)   # 各デバイスは 2 つの行を見ます
    return quantize(x)

qrows = quantize_shards(rows)
print(jax.typeof(qrows))
print(qrows.qvalue.sharding.spec, qrows.scale.sharding.spec)

q8[8@i,3]
P('i', None) P('i',)

そして、各デバイスが独自の行を持つシャードごとに QArray を見る場所での消費：

@jax.jit
@jax.shard_map(in_specs=QArrayP(jax.P('i')), out_specs=jax.P('i'))
def dequantize_shards(qx):
    assert jax.typeof(qx).shape == (2, 3)  # デバイスごとの QArray シャード
    return dequantize(qx)

print(jnp.max(jnp.abs(dequantize_shards(qrows) - rows)))

スケールが各行ごとのため、シャードごとに変換することはグローバルに変換することと完全に一致します——これがバッチ処理を快適にしました：

qrows_global = quantize(rows)
assert (qrows.qvalue == qrows_global.qvalue).all()
assert (qrows.scale == qrows_global.scale).all()

（

shard_map

to_tangent_spec

to_ct_spec

nospec

追加の例

レシピは常に同じです——値クラス、

lo_ty

lower_val

raise_val

HiType

ランク 1 アレイ

ランク 1 アレイは、2 つのベクトルの外積として m × n マトリックスを表し、

col : f32[m]

row : f32[n]

@dataclass(frozen=True)
class Rank1:
    col: jax.Array   # f32[m]
    row: jax.Array   # f32[n]

@dataclass(frozen=True)
class Rank1Ty(HiType):
    shape: tuple[int, int]   # (m, n)、密度で表された形状
    sharding: NamedSharding  # 密度の形状のシェーディング；因子はそこから導かれる

    def lo_ty(self):
        (m, n), spec = self.shape, self.sharding.spec
        return [ShapedArray((m,), jnp.dtype('float32'),
                            sharding=self.sharding.update(spec=jax.P(spec[0]))),
                ShapedArray((n,), jnp.dtype('float32'),
                            sharding=self.sharding.update(spec=jax.P(spec[1]))) ]
    
    def lower_val(self, r1):
        return [r1.col, r1.row]
    
    def raise_val(self, col, row):
        return Rank1(col, row)

    # ランク 1 マトリックスは加算の下で閉じていないため、接線型は密度です。
    def to_tangent_aval(self):
        return ShapedArray(self.shape, jnp.dtype('float32'),
                           sharding=self.sharding)

    def str_short(self, short_dtypes=False, mesh_axis_types=False):
        dims = [str(d) if p is None else f'{d}@{p}'
                for d, p in zip(self.shape, self.sharding.spec)]
        return f'r1[{",".join(dims)}]'
    __repr__ = str_short

def typeof_rank1(r1):
    col_s, row_s = jax.typeof(r1.col).sharding, jax.typeof(r1.row).sharding
    sharding = col_s.update(spec=jax.P(col_s.spec[0], row_s.spec[0]))
    return Rank1Ty((r1.col.shape[0], r1.row.shape[0]), sharding)

register_hitype(Rank1, typeof_rank1)

型はそれが表す密度形状を記録し、

r1[6,5]

QArrayTy

lo_ty

col

row

QArray

quantize

class Outer(VJPHiPrimitive):
    def __init__(self, col_aval, row_aval):
        if not (len(col_aval.shape) == 1 and len(row_aval.shape) == 1):
            raise TypeError(f'bad outer factor types: {col_aval}, {row_aval}')
        sharding = col_aval.sharding.update(
            spec=jax.P(col_aval.sharding.spec[0], row_aval.sharding.spec[0]))
        self.in_avals = (col_aval, row_aval)
        self.out_aval = Rank1Ty((col_aval.shape[0], row_aval.shape[0]), sharding)
        self.params = {}
        super().__init__()

    def expand(self, col, row):
        return Rank1(col, row)

    def vjp_fwd(self, nzs_in, col, row):
        return self(col, row), (col, row)

    def vjp_bwd_retval(self, res, g):   # g は密度、f32[m, n]
        col, row = res
        return g @ row, col @ g

class Factors(VJPHiPrimitive):
    def __init__(self, r1_aval):
        self.in_avals = (r1_aval,)
        self.out_aval = tuple(r1_aval.lo_ty())  # ファクター型は lo タイプです。
        self.params = {}
        super().__init__()

    def expand(self, r1):
        return (r1.col, r1.row)

class MatmulR1(VJPHiPrimitive):
    def __init__(self, x_aval, r1_aval):
        if not (isinstance(r1_aval, Rank1Ty) and len(x_aval.shape) == 2 and
                x_aval.shape[1] == r1_aval.shape[0]):
            raise TypeError(f'bad matmul_r1 operand types: {x_aval} @ {r1_aval}')
        self.in_avals = (x_aval, r1_aval)
        out_sharding = x_aval.sharding.update(
            spec=jax.P(x_aval.sharding.spec[0], r1_aval.sharding.spec[1]))
        self.out_aval = ShapedArray((x_aval.shape[0], r1_aval.shape[1]),
                                    jnp.dtype('float32'), sharding=out_sharding)
        self.params = {}
        super().__init__()

    def expand(self, x, r1):
        return jnp.outer(x @ r1.col, r1.row)   # 決して col ⊗ row を物質化しない

    def vjp_fwd(self, nzs_in, x, r1):
        col, row = factors(r1)  # 則は追跡されたコード：属性ではなくプライミティブ
        return self(x, r1), (x, col, row)

    def vjp_bwd_retval(self, res, g):
        x, col, row = res
        return jnp.outer(g @ row, col), x.T @ g

def outer(col, row):
    return Outer(jax.typeof(col), jax.typeof(row))(col, row)

def factors(r1):
    return Factors(jax.typeof(r1))(r1)

def matmul_r1(x, r1):
    return MatmulR1(jax.typeof(x), jax.typeof(r1))(x, r1)

（

Factors

in_avals

out_aval

MatmulR1

expand

col = jnp.arange(6., dtype='float32') / 6.
row = jnp.arange(5., dtype='float32') / 5.
r1 = outer(col, row)
print(jax.typeof(r1))

acts = jnp.ones((3, 6), 'float32')
print(jnp.max(jnp.abs(matmul_r1(acts, r1) - acts @ jnp.outer(col, row))))

r1[6,5]
0.0

Jaxprs、jit、そして自動微分はすべて以前と同じように機能します。ランク 1 演算子に関する勾配は密度であり、接線型が規定する通り、コンスストラクターを通過したチェーンされた勾配は因子に戻ります：

jax.jit(matmul_r1).trace(acts, r1).jaxpr

{ lambda ; a:f32[3,6] b:r1[6,5]. let
    c:f32[3,5] = call_hi_primitive[_prim=MatmulR1[{}]] a b
  in (c,) }

def r1_loss(x, r1):
    return jnp.sum(matmul_r1(x, r1) ** 2)

print(jax.typeof(jax.grad(r1_loss, argnums=1)(acts, r1)))

def factor_loss(col, row):
    return jnp.sum(matmul_r1(acts, outer(col, row)) ** 2)

g_col, g_row = jax.grad(factor_loss, argnums=(0, 1))(col, row)
print(jax.typeof(g_col), jax.typeof(g_row))

float32[6,5]
float32[6] float32[5]

そして、タイピング則は

QArrayTy

print(jax.typeof(outer(jax.device_put(jnp.zeros(8, 'float32'), jax.P('i')),
                       row)))

ここでは止まりましたが、vmap、scan、そして shard_map サポートは上記の節と同じレシピに従います：

dec_rank

inc_rank

leading_axis_spec

HiPspec

shard

unshard

タプル

私たちの最後の例は一般的なコンテナです——要素が任意の JAX 値であるタプルです。

QArrayTy

Rank1Ty

TupTy

import itertools

@dataclass(frozen=True)
class HiTup:
    elts: tuple

@dataclass(frozen=True)
class TupTy(HiType):
    tys: tuple  # コンポーネント型：アレイ型、または他の hi 型

    # lowering はコンポーネント型へ委任します。
    def lo_ty(self):
        return [lo for ty in self.tys for lo in ty.lo_ty()]
    
    def lower_val(self, tup):
        return [lo for ty, elt in zip(self.tys, tup.elts)
                for lo in ty.lower_val(elt)]
    
    def raise_val(self, *los):
        los = iter(los)
        return HiTup(tuple(ty.raise_val(*itertools.islice(los, len(ty.lo_ty())))
                          for ty in self.tys))

    # および接線型も同様に。
    def to_tangent_aval(self):
        return TupTy(tuple(ty.to_tangent_aval() for ty in self.tys))

    def str_short(self, short_dtypes=False, mesh_axis_types=False):
        return ('Tup{' +
                ','.join(t.str_short(short_dtypes, mesh_axis_types)
                         for t in self.tys) + '}')
    __repr__ = str_short

register_hitype(HiTup, lambda t: TupTy(tuple(map(jax.typeof, t.elts))))

委任によって無料に 2 つのものが得られます。第一に、要素自体が hi 型になり得ます——タプルのタプル、または QArray を持つタプル——なぜなら

lo_ty

QArrayTy

Rank1Ty

コンスストラクターおよびアクセス用プライミティブはランク 1 の例からよく知られています。新しい要素：要素インデックスは静的パラメータであり、

params

self.idx

@dataclass(frozen=True)
class TupSpec(MappingSpec):
    val: tuple  # コンポーネントごとに 1 つの軸エントリ

def tup_dec_rank(self, size, spec):
    return TupTy(tuple(ty.dec_rank(size, s)
                       for ty, s in zip(self.tys, spec.val)))

def tup_inc_rank(self, size, spec):
    return TupTy(tuple(ty.inc_rank(size, s)
                       for ty, s in zip(self.tys, spec.val)))

TupTy.dec_rank = tup_dec_rank
TupTy.inc_rank = tup_inc_rank

class MakeTup(VJPHiPrimitive):
    def __init__(self, elt_avals):
        self.in_avals = tuple(elt_avals)
        self.out_aval = TupTy(tuple(elt_avals))
        self.params = {}
        super().__init__()

    def expand(self, *elts):
        return HiTup(elts)

    def batch(self, axis_data, args, in_dims):
        return make_tup(*args), TupSpec(tuple(in_dims))

class GetTupElt(VJPHiPrimitive):
    def __init__(self, tup_aval, idx):
        self.in_avals = (tup_aval,)
        self.out_aval = tup_aval.tys[idx]
        self.params = dict(idx=idx)
        super().__init__()

    def expand(self, tup):
        return tup.elts[self.idx]

    def batch(self, axis_data, args, in_dims):
        tup, = args
        spec, = in_dims
        if spec is None:
            return get_tuple_element(tup, self.idx), None
        return get_tuple_element(tup, self.idx), spec.val[self.idx]

def make_tup(*elts):
    return MakeTup(map(jax.typeof, elts))(*elts)

def get_tuple_element(tup, idx):
    return GetTupElt(jax.typeof(tup), idx)(tup)

（

GetTupElt

in_dims

None

タプルが機能し、それらは入れ子になり、他の hi タイプを持っています：

tup = make_tup(jnp.arange(3.), 5.)
print(jax.typeof(tup))
print(get_tuple_element(tup, 1))

nested = make_tup(make_tup(1., 2.), jnp.arange(2.))
print(jax.typeof(nested))
print(jax.jit(lambda t: get_tuple_element(get_tuple_element(t, 0), 1))(nested))

print(jax.typeof(make_tup(qx, 3.)))  # 定量化アレイ要素

Tup{float32[3],float32[]}
5.0
Tup{Tup{float32[],float32[]},float32[2]}
2.0
Tup{q8[2,3],float32[]}

コンポーネントごとのマッピング仕様は vmap 下の得です：各要素に独自の

in_axes

out_axes

def swap(t):
    a, b = get_tuple_element(t, 0), get_tuple_element(t, 1)
    return make_tup(b, a)

out = jax.vmap(swap, in_axes=TupSpec((0, None)), out_axes=TupSpec((None, 0)),
               axis_size=3)(tup)
print(jax.typeof(tup), '->', jax.typeof(out))

Tup{float32[3],float32[]} -> Tup{float32[],float32[3]}

そして、コンポーネント型が独自のシェーディングを持つため、型の内部的シェーディングには追加の何ものも必要ありません：

print(jax.typeof(make_tup(rows, jnp.float32(1.))))

Tup{float32[8@i,3],float32[]}

これらのプライミティブ上の自動微分則、そして scan および shard_map サポート（コンポーネントごとの

HiPspec

tests/hijax_test.py

TupTy

私たちがカバーしなかったこと

プライミティブ側では、リマテリアライゼーションおよびデッドコード除去のカスタマイズのためのフックもあります。

Hijax ではいつも通り、

tests/hijax_test.py