このガイドでは、メリディアンで JAX バックエンドを使用する方法について説明します。
JAX バックエンドの概要
メリディアンでは、コアとなる数値演算と確率的なマルコフ連鎖モンテカルロ（MCMC）サンプリングのデフォルトの数値バックエンドとして JAX が使用されます（メリディアン 2.0 以降）。JAX は関数型プログラミング スタイルを推奨し、XLA（Accelerated Linear Algebra）コンパイルを利用して高度なパフォーマンス最適化とメモリ効率を実現します。
以前の TensorFlow バックエンドは非推奨となり、今後のリリースで削除される予定です。
チュートリアル: JAX の動作を確認するには、JAX スタートガイドのノートブックをご覧ください。
バックエンドの構成
デフォルトでは、メリディアンは JAX で実行されます。JAX を使用するために環境変数を構成する必要はありません。
以前の TensorFlow バックエンド（非推奨）
以前の TensorFlow バックエンドで一時的に実行する必要がある場合は、メリディアンをインポートする前に MERIDIAN_BACKEND 環境変数を 'tensorflow' に設定します。
import os
# Select legacy TensorFlow backend (deprecated)
os.environ['MERIDIAN_BACKEND'] = 'tensorflow'
# Now it is safe to import Meridian modules
from meridian.model import model
from meridian.data import load
精度設定
デフォルトでは、メリディアンは JAX で 64 ビット精度（float64）で実行されます。
次のような目的がある場合、ユーザーは 32 ビット精度（float32）での実行を代わりに選択できます。
- トレーニング ランタイムの高速化: 32 ビット浮動小数点演算は、ハードウェア アクセラレータ（GPU や TPU など）で高速に実行できます。
- メモリ使用量の削減: 32 ビット精度により、MCMC サンプリング中のメモリ消費量が削減されます。
32 ビット精度を使用するには、メリディアンをインポートする前に MERIDIAN_ENABLE_JAX_X64 環境変数を 'False'（または '0'）に設定します。
import os
# Disable 64-bit precision (enable 32-bit precision)
os.environ['MERIDIAN_ENABLE_JAX_X64'] = 'False'
# Now it is safe to import Meridian modules
from meridian.model import model
MERIDIAN_ENABLE_JAX_X64 環境変数が設定されていない場合、または 'True' あるいは '1' に設定されている場合、メリディアンはデフォルトで 64 ビット精度を使用します。
型の整合性
メリディアンはデフォルトで 64 ビット精度で動作するため、ユーザー指定の値、カスタム配列、分布パラメータのすべてで型の整合性が維持されるようにしてください。
- 浮動小数点数リテラルと配列: 標準の Python 浮動小数点数リテラル（
0.2、0.9など）は、デフォルトで 64 ビット浮動小数点数になります。事前分布または入力用の NumPy 配列を作成する場合は、np.float64またはdtype=np.float64を使用してデフォルトの精度に合わせます。 - 精度を一致させてカスタム事前分布を構築する:
PriorDistributionでカスタム事前分布を定義する場合は、すべての分布パラメータ（loc、scale、concentration0、concentration1など）が、適用されている精度（デフォルトでは 64 ビット浮動小数点数）と一致するようにします。
JAX と TensorFlow における API の違い
JAX バックエンドを使用する際、以下の API の主な違いに注意する必要があります。
事前分布
メリディアン モデルは、JAX（tensorflow_probability.substrates.jax）で TensorFlow Probability を使用します。JAX でカスタム事前分布を構成する場合は、tensorflow_probability.substrates.jax as tfp_jax をインポートし、tfp_jax.distributions を使用して分布を構築します。
すべてのカスタム分布パラメータで 64 ビット精度（np.float64 や 64 ビット浮動小数点数配列など）を使用し、メリディアンのデフォルトの精度設定との型の整合性を維持します。
JAX
import numpy as np
import tensorflow_probability.substrates.jax as tfp_jax
from meridian.model import constants
from meridian.model import prior_distribution
# Parameters use 64-bit precision
roi_mu = np.float64(0.2)
roi_sigma = np.float64(0.9)
prior = prior_distribution.PriorDistribution(
roi_m=tfp_jax.distributions.LogNormal(
roi_mu, roi_sigma, name=constants.ROI_M
)
)
TensorFlow（非推奨）
import tensorflow_probability as tfp
from meridian.model import constants
from meridian.model import prior_distribution
roi_mu = 0.2
roi_sigma = 0.9
prior = prior_distribution.PriorDistribution(
roi_m=tfp.distributions.LogNormal(
roi_mu, roi_sigma, name=constants.ROI_M
)
)
明示的なシードが必要
JAX バックエンドを使用する場合、確率的関数（sample_posterior() など）には明示的なシードが必要です。TensorFlow では、ランダムシードを自動的に選択するグローバルな乱数ジェネレータが使用されますが、JAX では、このシードを明示的に指定する必要があります。シード値の違いによって、費用対効果の推定値や予算配分に統計的に有意な差が生じることはありません。
# Explicitly set a seed for MCMC sampling when using the JAX backend
mmm.sample_posterior(
n_chains=2,
n_adapt=1000,
n_burnin=500,
n_keep=1000,
seed=0,
)
JAX の乱数とシードの詳細については、JAX の疑似乱数に関するドキュメントをご覧ください。
数値の違いと再現性
TensorFlow と JAX では計算グラフのコンパイル方法が異なるため、同じデータと乱数シードを使用して JAX に切り替えると、事後推定値にわずかな数値の違いが生じる場合があります。
バックエンド間で事後分布が同一でない場合もありますが、その差は通常小さく、費用対効果や予算配分などのビジネス指標では統計的に有意ではありません。これにより、JAX バックエンドへの切り替え後も、モデルの分析情報の整合性が維持されます。
パフォーマンスに関する注意事項
内部テストでは、JAX によって初期モデルの実行速度が大幅に上昇し、GPU 使用時の TensorFlow と比較して、平均実行時間が約 40%、メモリ使用量が約 70% 削減されることがわかりました。また、JAX はモデルのイテレーションも効率化することで、実行時間を 2 倍高速化してメモリ使用量を 4 分の 1 に削減し、カーネルの再起動を不要にしてワークフローの中断をなくしました。
メモリ効率が向上したことで、計算負荷の高いパラメータを調整する余地が広がりました。たとえば、Meridian.sample_posterior() では、unrolled_leapfrog_steps 引数を引き上げることが可能です（例: 1 を 5 に）。これにより、ハードウェアのメモリ上限を超えずに No-U-Turn-Sampler（NUTS）の軌跡を長くして、収束を加速できます。また、適応フェーズでの収束をさらに促進できるよう、n_adapt パラメータの値を引き上げることも可能です。