hijax で JAX に新しい型を定義する

Defining new Jax types with hijax

JAX の組み込みの通貨は配列であり、変換される関数は配列を入力として受け取り配列を出力します。しかし、集約データを扱う場合、通常は pytree を使用しますが、透明性が望ましくない場合があります。hijax 型（または「hi 型」）は、独自のアイデンティティを持つ新しい型を定義する方法を提供します。このドキュメントでは、量子化配列型を例に、HiType のサブクラス化、値クラスの登録、hijax プリミティブの記述方法を説明します。これにより、jaxpr 内で単一の値として扱われ、内部不変条件を強制し、カスタムのタンジェント型、vmap バッチング、シャーディング情報を型に持たせることができます。