izmyonの日記

Preface この記事では、JAX学習記録の第三弾として、JAXにおける自動微分(jax.grad)および関数の自動ベクトル化(jax.vmap)についてまとめる。まず、jax.gradとセットで使われることが多いAutomatic Vectorization (jax.vmap)について説明する。jax.vmapは、jax.gradやjax.jitなどと並び、JAXの大きな特徴の一つである「関…