TensorFlow Advent Calendar 2020 10日目の記事です。空いてたので当日飛び入りで参加しました。 この記事では、TensorFlowの関連ライブラリである「JAX」について初歩的な使い方、ハマりどころ、GPU・TPUでの使い方や、画像処理への応用について解説します。 JAXとは https://github.com/google/jax Google製のライブラリで、AutogradとXLAからなる、機械学習のための数値計算ライブラリ。簡単に言うと「自動微分に特化した、GPUやTPUに対応した高速なNumPy」。NumPyとほとんど同じ感覚で書くことができます。自動微分については解説が多いので、この記事では単なる高速なNumPyの部分を中心に書いていきます。 関連記事 JAX Quickstart JAXで始めるディープラーニング JAX : Tutorials