jaxtypingの紹介

テンソルの形状とデータ型に対しても型チェックを行えるjaxtypingというpythonライブラリを紹介しています。

Transcript

1. AI Community 2025.10.02 小林 茂樹 GOドライブ株式会社 jaxtypingの紹介
2. 2 項目 01｜モチベーション 02｜型アノテーションの書き方 03｜実行時型チェック 04｜まとめと所感
3. 3 01 モチベーション
4. 4 ▪ NumpyやPyTorchでテンソルを扱う際に次元（shape）が コロコロ変わる ▪ コードを読んだだけではshapeが分からず、具体的な処理 が掴めないことがある ▪ 変数名や処理の背景からshapeを推測しながらコードを読 むことが多いかと思います ▪ 例：torch.tensorの画像なら (B C H W) だろうと推測する shapeの推測には限界がある shapeが分からない辛さ
5. 5 1. poseという変数がある 2. 位置姿勢があると推測 3. 姿勢の表現方法によってshapeが違う a. オイラーのshape：(3,) b. クォータ二オンのshape：(4,) c. 回転行列のshape：(3,3) shapeの推測が困難な例 変数名だけではshapeの推測が困難
6. 6 以下のようなshapeに対してのコメントやアサーションを書いたことは ありませんか？ shapeが分からない辛さ
7. 7 shapeに対してコメントを書いても良いが、ちゃんと型 チェックしたい jaxtypingのススメ jaxtypingを使いましょう！
8. 8 shape, dtypeに対しても型アノテーションを行うライブラリ ▪ リポジトリ：https://github.com/patrick-kidger/jaxtyping ▪ ドキュメント：https://docs.kidger.site/jaxtyping/ ▪ 名前的にJAXしか対応していなさそうだが、 PyTorch/NumPy/TensorFlowにも対応 ▪ 実行時型チェックもできる jaxtyping
9. 9 02 型アノテーションの書き方
10. 10 例 https://docs.kidger.site/jaxtyping/#example
11. 11 基本形 dtype[array, shape] 基本は以下のような型になる 例えばPyTorchで画像を扱う場合
12. 12 array 基本は以下のような型になる 例えばPyTorchで画像を扱う場合 dtype[array, shape]
13. 13 array 複数のライブラリのarrayを使用することが可能 ▪ jax.Array ▪ np.ndarray ▪ torch.Tensor ▪ tf.Tensor ▪ mx.array https://docs.kidger.site/jaxtyping/api/array/#array
14. 14 shape 基本は以下のような型になる 例えばPyTorchで画像を扱う場合 dtype[array, shape]
15. 15 shapeの書き方 str型でshapeを指定した場合、変数となる int型で指定することで定数であることを明示できる 3チャネルの224x224の画像の場合 https://docs.kidger.site/jaxtyping/api/array/#shape これらの変数は関数内で一貫した値にならなければエラーになる
16. 16 shapeの書き方 ローカル・インスタンス変数などの値を次元数として使う こともできる インスタンス変数の値を 次元数として設定 https://docs.kidger.site/jaxtyping/api/array/#shape
17. 17 shapeの書き方 変数を用いた計算からshapeを設定できる dimから1引いた次元数 をshapeに指定 https://docs.kidger.site/jaxtyping/api/array/#shape
18. 18 shapeの書き方 アンダースコアを設定すれば該当の次元は型チェックを避け られる チャネルの次元は 型チェックされない https://docs.kidger.site/jaxtyping/api/array/#shape
19. 19 shapeの書き方 =を使うとコメントとなる（型チェックはされない） shapeに対して型チェック は行われない https://docs.kidger.site/jaxtyping/api/array/#shape
20. 20 shapeの書き方 *を付けた軸は連続する0個以上の軸にできる shapeが (N B C H W) の場合、以下のように書ける https://docs.kidger.site/jaxtyping/api/array/#shape
21. 21 shapeの書き方 ... を0個以上の軸とすることができる shapeが (N B C H W) の場合、以下のように書ける https://docs.kidger.site/jaxtyping/api/array/#shape
22. 22 dtype 基本は以下のような型になる 例えばPyTorchで画像を扱う場合 dtype[array, shape]
23. 23 dtype 基本的なデータ型が用意されている ▪ Bool ▪ Float ▪ Int ▪ UInt https://docs.kidger.site/jaxtyping/api/array/#dtype
24. 24 dtype ▪ Float16 ▪ Float32 ▪ Float64 特定の精度のデータ型もある ▪ Int2 ▪ Int4 ▪ Int8 ▪ Int16 ▪ Int32 ▪ Int64 ▪ UInt2 ▪ UInt4 ▪ UInt8 ▪ UInt16 ▪ UInt32 ▪ UInt64 ▪ BFloat16 https://docs.kidger.site/jaxtyping/api/array/#dtype
25. 25 03 実行時型チェック
26. 26 実行時型チェックを呼び出すhook3種類 ▪ 関数単位でのhook ▪ モジュール単位でのhook ▪ pytestに対してのhook hookの種類 https://docs.kidger.site/jaxtyping/api/runtime-type-checking/
27. 27 実行時型チェックを行いたい関数に対してjaxtypedというデコレータを設定する データクラスの場合は__init__の中に対して実行時型チェックが行われる 関数単位でのhook https://docs.kidger.site/jaxtyping/api/runtime-type-checking/#jaxtyping.jaxtyped
28. 28 install_import_hookのwithブロック内でimportする このブロック内でimportされたモジュールの全関数に対して実行時型チェックが 適用される モジュール単位でのhook https://docs.kidger.site/jaxtyping/api/runtime-type-checking/#jaxtyping.install_import_hook
29. 29 pytestに追加で引数を与えれば良い jaxtyping-packagesに実行時型チェックを適用したいモジュールを指定する 最もキレイにまとまるが、テストに対してしか実行時型チェックができない pytestに対してのhook pyproject.tomlで設定する場合 pytest.iniで設定する場合 https://docs.kidger.site/jaxtyping/api/runtime-type-checking/#pytest-hook
30. 30 pytorch-lightningを用いている場合はfast_dev_runというパラメーターを設定 することでtrainingとvalidationを1度だけ実行することができる ➡ ユニットテストでtraining、validationに対してのjaxtypingの型チェックを少 ない計算量で行うことができる （余談）pytorch-lightningを用いたテスト https://lightning.ai/docs/pytorch/stable/common/trainer.html#fast-dev-run
31. 31 04 まとめと所感
32. 32 IDEのホバー情報にshape、dtypeに対しての型情報が出てこない 改善されると嬉しいところ shape、dtypeについて の情報がない
33. 33 改善されると嬉しいところ ログが長い ▪ テンソルの内容が全てログに出てきて本質的なログが分かりづらい サンプルコード エラーログ 想定している返り値の shapeではないので エラーになる エラー毎にテンソルの 生値が出てくるので、 エラーが冗長になる
34. 34 NumpyやPyTorchなどのarrayに対してshapeとdtypeも含めて型アノテーション と型チェックができるライブラリ ▪ jaxtypingを導入したリポジトリを約1年ぶりに触ったが、可読性が上がって おり理解を早めた（と思う） ▪ コードを見ただけでshapeが分かるのはかなり嬉しい ▪ もっと流行ってほしい まとめと所感