An Introduction to JAX

Abstract

In this talk, I provide an introduction to JAX and Flax as a Deep Learning framework. I will discuss the basic tools that a machine learning framework has to provide, and show step by step how it is implemented in JAX with Flax. Finally, a short practical demo is given based on our Deep Learning Tutorials.