Penzai - ニューラルネットワークの構築・編集・可視化のためのJAXツールキット

• Google DeepMindが公開した、読みやすく関数型のPytreeデータ構造でモデルを記述するためのJAXライブラリ
• モデルの学習後でも容易に扱えるようにすることに重点を置いている
  • モデル構成要素のリバースエンジニアリングや削除、内部活性の検査・調査、モデル手術、アーキテクチャのデバッグなどに関する研究に適している
  • モデルの可視化、修正、分析のためのツールを提供しており、単にモデルを構築して学習させる用途にも利用できる
• 一緒に設計されているが、それぞれ独立して使用できるモジュール型ツール群で構成

含まれるツール

penzai.nn (pz.nn)

• Flax、Haiku、Keras、Equinoxなどの他のニューラルネットワークライブラリの代替となる、宣言的コンビネータベースのニューラルネットワークライブラリ
• モデルのforward pass全体の構造をモデルpytreeに公開し、pretty printingによってモデルのすべての処理を確認でき、jax.tree_utilで新しいランタイムロジックを注入できる
• モデルは内部的には呼び出し可能なpytreeにすぎない

penzai.treescope (pz.ts)

• 通常のIPython/Colabレンダラーを置き換えるドロップイン方式で動作する、強力な対話型Python pretty-printer
• Penzaiモデルやその他の深くネストされたJAX pytreeを理解しやすくするよう設計されており、任意次元のNDArray可視化を標準でサポートする

penzai.core.selectors (pz.select)

• JAXの.at[...].set(...)構文を任意の型ベースのpytree走査へ一般化する、pytree向けの多機能ツール
• Penzaiモデルやその他のデータ構造に対する複雑な書き換えやオンザフライのパッチを容易に実行できる

penzai.core.named_axes (pz.nx)

• 通常のJAX関数を名前付き軸に対してベクトル化し、新しい配列APIを学ばなくても名前付きプログラミングスタイルと位置ベースのプログラミングスタイルをシームレスに切り替えられる軽量な名前付き軸システム

penzai.data_effects (pz.de)

• pytree走査を基盤として構築され、モデルの記述や利用を妨げることなく制御を可能にする、副次引数、乱数、状態変数のためのオプトインシステム

[brainer]

PyTorch向けにも似たようなものはあるのでしょうか？