Preface この記事ではJAX学習記録シリーズ第四弾として、JAXの自動微分がどのように実現されているか、詳細を見てみる。正直なところ、特に数学的な説明に関して理解しきれていないところが多く、基本的には原文に忠実に訳しており、自分の理解による説明は…

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

Preface この記事ではJAX学習記録シリーズ第二弾として、Pytreeについてまとめる。基本的に第一弾は読んでる前提で書くので、もし読んでない方がいたら読むことをお勧めする。Pytreeは以下の第一弾のブログでも出てきたように、モデルパラメータを格納するデ…

Preface このシリーズは、JAXのドキュメントを読みながら学習した際に、ノート代わりに記事に残したものである。JAXを理解する際の要所は、XLAやjust-in-time、Asynchronous dispatchなど多岐にわたるが、この記事ではまず第一弾として、Numpy-likeなAPIで利…