PyTensor，JAX，Flax のインストール（Ubuntu 上）

PyTensor は，多次元配列を含む数式を定義・最適化・コンパイルして実行する Python のライブラリである。Theano の後継として PyMC 開発チームが開発を継続しているものであり，Theano と同系統の記号的微分と式コンパイルの機能を持つ。

JAX は，NumPy 互換の API に自動微分（grad）と JIT コンパイル（jit），ベクトル化（vmap），CPU・GPU・TPU での実行を組み合わせた数値計算ライブラリである。Flax は JAX 上でニューラルネットワークを構築・学習する軽量ライブラリであり，Lasagne が Theano に対して果たしていた役割に相当する。現在の API は Flax NNX である。Optax は JAX 用の最適化アルゴリズムのライブラリである。

Theano の開発は 2017 年に終了しており，最終版は 2020 年公開の 1.0.5 である。Theano 上のニューラルネットワークライブラリである Lasagne も 2015 年公開の 0.1 が最後の公式リリースである。いずれも NumPy 2.x や近年の Python では動作しないため，このページでは後継および同等の機能を持つ PyTensor，JAX，Flax のインストールを説明する。

前準備

Ubuntu のシステム更新

Ubuntu で OS のシステム更新を行うときは，端末で，次のコマンドを実行する。これは，パッケージ情報を最新の状態に保ち，インストール済みのパッケージをセキュリティ更新やバグ修正を含めて更新するためである。

Ubuntu のインストールはこちらの別ページで説明する。

# パッケージリストの情報を更新
sudo apt update
# インストール済みのパッケージを包括的に更新 (依存関係も考慮)
sudo apt full-upgrade
# カーネル更新等で実際に再起動が必要な場合のみ実行
# sudo shutdown -r now

前提ソフトウェアのインストール

PyTensor は，式を C コードに変換して実行するため C/C++ コンパイラと BLAS の実装を利用する。端末で，次のコマンドを実行する。

# パッケージリストの情報を更新
sudo apt update
sudo apt -y install build-essential gcc g++ libopenblas-dev liblapack-dev
sudo apt -y install git wget python3 python3-dev python3-venv python3-pip

Python の仮想環境の作成

Ubuntu 23.04 以降では，OS が管理する Python に pip で直接パッケージをインストールすることは PEP 668 により禁止されており，externally-managed-environment エラーとなる。仮想環境を作成し，その中にインストールする。端末で，次のコマンドを実行する。

python3 -m venv ~/jaxenv
source ~/jaxenv/bin/activate
pip install -U pip setuptools wheel

端末を開き直したときは，端末で，次のコマンドを実行して仮想環境を有効にする。

source ~/jaxenv/bin/activate

PyTensor のインストール

公式のインストール手順は次のページで公開されている。

https://pytensor.readthedocs.io/

PyTensor のインストール
仮想環境を有効にした状態で，端末で，次のコマンドを実行する。
```
source ~/jaxenv/bin/activate
pip install -U pytensor
```
動作テスト
式の定義，コンパイル，記号的微分が動作することを確認する。仮想環境を有効にした状態で，端末で，次のコマンドを実行する。
```
source ~/jaxenv/bin/activate
python -c "
import pytensor
import pytensor.tensor as pt
x = pt.dscalar('x')
y = x ** 2
f = pytensor.function([x], [y, pt.grad(y, x)])
print(pytensor.__version__)
print(f(3.0))
"
```
バージョン番号と [array(9.), array(6.)] が表示されれば，インストールは成功している。

JAX，Flax のインストール

公式のインストール手順は次のページで公開されている。

https://docs.jax.dev/en/latest/installation.html

https://flax.readthedocs.io/

JAX のインストール
CPU のみで実行する場合は，仮想環境を有効にした状態で，端末で，次のコマンドを実行する。
```
source ~/jaxenv/bin/activate
pip install -U jax
```
NVIDIA GPU（CUDA 12）を使用する場合は，仮想環境を有効にした状態で，端末で，次のコマンドを実行する。CUDA のライブラリは pip で同時にインストールされるため，別途 CUDA Toolkit をインストールする必要はない。NVIDIA ドライバは 525 以降が必要である。
```
source ~/jaxenv/bin/activate
pip install -U "jax[cuda12]"
```
Flax，Optax のインストール
仮想環境を有効にした状態で，端末で，次のコマンドを実行する。
```
source ~/jaxenv/bin/activate
pip install -U flax optax
```
JAX の動作テスト
自動微分と実行デバイスを確認する。仮想環境を有効にした状態で，端末で，次のコマンドを実行する。
```
source ~/jaxenv/bin/activate
python -c "
import jax
import jax.numpy as jnp
print(jax.__version__)
print(jax.devices())
print(jax.grad(lambda x: x ** 2)(3.0))
"
```
バージョン番号，利用可能なデバイスの一覧，6.0 が表示されれば，インストールは成功している。GPU を使用する設定でインストールした場合は，デバイスの一覧に CudaDevice が表示される。

Flax の動作テスト

多層パーセプトロンを定義し，学習を 1 ステップ実行して確認する。仮想環境を有効にした状態で，端末で，次のコマンドを実行する。

mkdir -p ~/code/flaxtest
cd ~/code/flaxtest
cat > mlp.py <<'EOF'
import jax
import jax.numpy as jnp
import optax
from flax import nnx


class MLP(nnx.Module):
    def __init__(self, din, dmid, dout, *, rngs):
        self.linear1 = nnx.Linear(din, dmid, rngs=rngs)
        self.linear2 = nnx.Linear(dmid, dout, rngs=rngs)

    def __call__(self, x):
        return self.linear2(nnx.relu(self.linear1(x)))


model = MLP(4, 16, 3, rngs=nnx.Rngs(0))
optimizer = nnx.Optimizer(model, optax.adam(1e-3), wrt=nnx.Param)

key = jax.random.key(0)
x = jax.random.normal(key, (32, 4))
y = jnp.zeros((32,), dtype=jnp.int32)


def loss_fn(model):
    logits = model(x)
    return optax.softmax_cross_entropy_with_integer_labels(logits, y).mean()


loss, grads = nnx.value_and_grad(loss_fn)(model)
optimizer.update(model, grads)
print("loss:", loss)
EOF

続けて，仮想環境を有効にした状態で，端末で，次のコマンドを実行する。

source ~/jaxenv/bin/activate
cd ~/code/flaxtest
python mlp.py

損失の値が表示されれば，インストールは成功している。

MNIST のサンプルプログラム
手書き数字認識（MNIST）を用いた学習の例は，Flax の公式チュートリアルで公開されている。

https://flax.readthedocs.io/en/latest/mnist_tutorial.html