Skip to content

A pathway and collection of resources to learning Jax from beginning to advance.

License

Notifications You must be signed in to change notification settings

deterministic-algorithms-lab/Jax-Journey

Repository files navigation

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 : jax-ml/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

About

A pathway and collection of resources to learning Jax from beginning to advance.

Topics

Resources

License

Stars

Watchers

Forks

Releases

No releases published

Packages

No packages published