hijax を使った新しい Jax 型の定義

hijax を使った新しい Jax 型の定義# JAX の基本的な単位は配列です。変換される関数は配列を受け取り、配列を返します。また、トレーシング機構が見るすべての中間値は f32[3,4] のような配列型を持ちます。集計データを扱いたい場合、通常のツールは pytree です。配列をコンテナにバンドルすると、JAX はそのバンドルを配列のリーフに透過的にフラット化します。しかし、時には透過性が望ましくないこともあります。一部のデータは、独自のアイデンティティを持つ新しい型としてモデル化するのが最善です。それは jaxprs では配列リーフの散布ではなく、単一の型の単一の値として表示されるべきです。内部的な不変量を持つため、ユーザーは固定された操作セットを通じてのみそれを生成・消費すべきであり、コンポーネントを自由に構築したりパターンマッチしたりするのではなく、その接線型は元の構造とは異なる場合があります。これにより、それに対する導関数は単に「接線用の同じ pytree」ではなくなります。また、vmap の下での独自のバッチ処理の概念を持つこともあります。型にシャーディング情報を含めることができ、JAX の明示的なシャーディングモードに参加できます。hijax 型（または「hi 型」）はこれを提供します。HiType をサブクラス化して型を定義し、Python クラスをその型の値を持つものとして登録し、新しい型を言及する入力と出力型を持つ hijax プリミティブを記述します。このドキュメントでは、量子化配列型という 1 つの実行例で、この全体像を説明します。hijax プリミティブにある程度の習熟があることを前提とします。それらの紹介については、hijax プリミティブによるカスタム導関数ルールを参照してください。hijax のすべてと同様に、これは実験的なものです。jax.experimental.hijax からのインポートを期待し、API が進化することを期待してください。TL;DR# HiType をサブクラス化し、lo_ty、lower_val、raise_val を実装して、型とその値が通常の（「lojax」）配列にどのように低下するかを説明します。次に、register_hitype を呼び出して、値クラスを型に関連付けます。新しい型を言及する VJPHiPrimitive サブクラスを記述します。これらは、型の値が生成および消費される唯一の方法です。自動微分の場合、型で to_tangent_aval を実装し、プリミティブで VJP/JVP ルールを実装します。vmap の場合、型で dec_rank と inc_rank を実装し、独自の設計のマッピング仕様サブクラスを実装し、プリミティブでバッチルールを実装します。マップされた hi 型引数には、明示的な axis_size と spec 値の in_axes/out_axes エントリが必要です。型でのシャーディング（明示モード）の場合、型にシャーディングデータを記録し（例: NamedSharding フィールド）、lo_ty で消費し、プリミティブのタイピングルールで伝播します。例：量子化配列# int8 に量子化された配列を扱いたいとします。量子化配列は実際には配列のペアです。int8 値と、各行で共有される浮動小数点スケールです（つまり、最後の軸に沿って、行ごとに 1 つのスケールで量子化します。一般的な行ごと/チャンネルごとの量子化スキームのように）。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 として登録して完了することもできます。しかし、失うものを考えてみましょう。不変量。2 つのコンポーネントは連動しています。scale は qvalue の最後の軸を除いた形状を持つ必要があり、qvalue は scale と一緒にのみ意味があります。pytree として、コードがストリームを横断するのを妨げるものはありません。変換の下では、JAX 自体は無関係なリーフのみを見ます。jaxprs の型。pytree として、量子化配列はトレースされたコードでは 2 つの無関係な配列値として表示されます。型付きの値が 1 つある方が、jaxprs が意味をなすでしょう。接線。量子化配列の値は離散グリッド上に存在するため、グリッドに沿ってそれらを摂動することは意味がありません。しかし、pytree の接線型は、リーフの接線型の pytree に強制されます。そして、整数配列（qvalue など）の接線型は float0 配列であり、取るに足らないペイロードしか運べません。したがって、pytree として、量子化配列はまったく有用な摂動を許容しないでしょう。代わりに、量子化配列全体に対して、量子化値が近似する連続的な f32 配列のような接線型を選択します。そのため、代わりに QArray を hijax 型にします。型# hijax 型は HiType のサブクラスです。必須のコアは小さいです。lo_ty は、どの lojax（配列）型がこの型を構成するかを示します。lower_val と raise_val は、値とその配列リストを相互に変換します。型はハッシュ可能で等価比較可能である必要があります（frozen dataclass が両方を提供します）。これは pytree の flatten/unflatten インターフェースに似ていますが、型のレベルにあります。型のみが与えられれば、JAX は値を用意せずに、低下した型を計算できます。また、型にシャーディングフィールドを与え、値がデバイス間でどのようにパーティション化されるかを記録します。これは、ドキュメントの最後のシャーディングセクションまで無視できます。qvalue コンポーネントのパーティショニングを記述するために JAX の NamedSharding を再利用し、最後の軸を削除することで scale のパーティショニングをそこから導出します。フィールドの特別な点はありません。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 のシャーディングはそこから導出されます。 # 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) # 自動微分：量子化配列の接線は通常の float 配列です（以下参照） def to_tangent_aval(self): return ShapedArray(self.shape, jnp.dtype('float32'), sharding=self.sharding) # printing, e.g. in 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 呼び出しは、値クラスを型に関連付けます。その 2 番目の引数は、任意の値を計算します。これは、jax.typeof が配列を ShapedArray 型にマッピングする方法に似ています。（私たちのものは、qvalue コンポーネントから形状とシャーディングの両方を読み取ります。すべての配列はシャーディングを運びます。メッシュがない場合は自明です。）実際、登録後、jax.typeof は QArrays で機能し、JAX 変換は値が期待される場所ならどこでもそれらを受け入れます。プリミティブ# pytree を使用すると、ユーザーは値を自由に構築および分解できます。hijax 型を使用すると、値は新しい型を言及する宣言された型を持つ hijax プリミティブによってのみ生成および消費されます。そこで不変量が強制されます。すべてのプリミティブがそれを維持すれば、それは常に保持されます。私たちの 2 つのプリミティブは quantize と dequantize で、Custom derivative rules with hijax primitives の VJPHiPrimitive API で記述されています。それぞれが入力と出力の型を宣言し、expand で実装を提供し、（自動微分を見越して）ストレートスルー推定器 VJP ルールを運びます。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(