はじめに

複雑な関数の微分を求めたい時に、自動微分が便利ですよね。そんな時、私は今までPytorchの自動微分を利用していました。しかし、自動微分だけ使いたいのにpytorchのパッケージがかなり重いので、もっと軽量なパッケージはないかなと探していたところ、JAXにたどりつきました。

JAXとは？

公式

Autograd(今はメンテナンスされていない)の更新バージョンです。GPUを使って高速に自動微分を計算できます(もちろんCPUでも動きます)。

インストール方法

CPU only version

pip install --upgrade pip
pip install --upgrade jax jaxlib  # CPU-only version

GPUを使いたい場合は公式のpip installationガイダンスを参照ください。

JAXでの二階微分

それでは試しにlog関数の二階微分を求めてみます。

import jax.numpy as jnp
from jax import grad

# log関数の定義
fn = lambda x0: jnp.log(x0)

# x = 1の周りで微分
x = 1

# 代入
y0 = fn(x)
# 一回微分
y1 = grad(fn)(x)
# 二階微分
y2 = grad(grad(fn))(x)

実行結果

>>> float(y0), float(y1), float(y2)
(0.0, 1.0, -1.0)

最後に

JAXを使うとお手軽簡単に自動微分が使えます。