#3 JAX Course - Future of ML research in JAX/Flax

[bhutanisanyam1]

Here are the link from prev sessions @kaczmarj:

@bhutanisanyam1

As promised, link to the colab being discussed today: Google Colab

@bbischof-wandb

Here are some of the relevant resources from today’s Lecture:

• Here is a link to the AutoDiff Notebook from Akash: TF_JAX_Tutorials - Part 9 (Autodiff in JAX) | Kaggle
• Here’s the PyTree Notebook from Akash: TF_JAX_tutorials - Part 10 (Pytrees in JAX) | Kaggle
• Here’s the link to “Flax Basics” from the Flax docs: Google Colab
• Here’s the link to “Annotated MNIST” from the Flax docs: Google Colab
• Here’s a link to the UNET lecture from last month: PyTorch Book Reading - 8. U-net, Image Segmentation and Image Augmentations in PyTorch - YouTube
• Here’s the Scenic implementation of U-net: scenic/unet.py at main · google-research/scenic · GitHub