playjax

Examples of JAX/NumPyro for a note on JAX/NumPyro written in Japanese:

機械学習で楽しむ JAX と NumPyro

contents

第 1 章 イントロダクション 2
第 2 章 JAX の自動微分で遊ぶ 3
2.1 自動微分 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 3
2.2 for ループを自動微分する . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 10
2.3 JIT とクラス . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 15
2.4 自動微分のカスタマイズ . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 16
2.5 jax.numpy のインデックス操作 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 18
第 3 章 JAX で最適化 20
3.1 JAX の optimizer で最適化を楽しむ . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 20
第 4 章 Numpyro の HMC-NUTS でマルコフ鎖モンテカルロ・シミュレーションをする 24
4.1 曲線フィット . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 24
4.2 カスタマイズした自動微分で numpyro を動かす . . . . . . . . . . . . . . . . . . . . . . . . 34
第 5 章 Numpyro で楽しむガウス過程 36
5.1 ガウス過程 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 36
5.2 ガウス過程ノイズを含んだモデルフィット . . . . . . . . . . . . . . . . . . . . . . . . . . . 45
第 6 章 NumPyro 47
6.1 scipy/JAX/NumPyro でサンプリング . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 47