Getting started with JAX

Powering the future of high-performance numerical computing and ML research

Photo by Lance Asper on Unsplash

Introduction

JAX is a Python library developed by Google to perform high-performance numerical computing on any type of device (CPU, GPU, TPU, etc…). One of the main applications of JAX is Machine Learning and Deep Learning research development, although the…