Jax-Journey

A first Introduction

Blog : Getting Started with Jax by Robert Tjarko Lange

Notebook : jax_basic.ipynb

Functional Programmning Intro

Blog : Functional JS 2 , Functional JS 3 by Krzysztof Czernek

How it works

A nice comment on Jax : google/jax#196 (comment)

Blog : From PyTorch to Jax : Towards NN frameworks that purify stateful code by Sabrina J. Mielke

Notebook : making_functions_pure.ipynb

Flax Intro

Blog : Documentation Tutorial

Notebook : Basics-MLP , Basics-MNIST

Note : Now is the best time to read Jax-The Sharp Bits

Haiku Intro

Notebook : haiku basics with lstm

Blog : From PyTorch to Jax : Towards NN frameworks that purify stateful code by Sabrina J. Mielke

Basic Transformer Using Haiku

Notebook : basic_transformer.ipynb

Blog : FineTuning Transformers with Jax+Haiku by Madison May

Jax - A deeper Dive

The following tutorials form a nice starting point, for a deeper dive into Jax mechanics :

1. On Automatic Differentiation in Jax

2. How Jax Represents your functions internally : Jaxprs

3. How Jax Primitives Work