jax-in-action

JAX in Action workshop files

• Colab: https://colab.research.google.com/github/yufengg/jax-in-action/blob/master/JAX_in_Action_part_1.ipynb

• Part 2: https://colab.research.google.com/github/jax-ml/jax-ai-stack/blob/main/docs/source/getting_started_with_jax_for_AI.ipynb