# \#2 JAX course - Working with Neural Networks in JAX

**URL:** <https://community.wandb.ai/t/2-jax-course-working-with-neural-networks-in-jax/1455>\
**Category:** JAX\
**Created:** [December 2, 2021, 1:47pm UTC](https://community.wandb.ai/t/2-jax-course-working-with-neural-networks-in-jax/1455 "2021-12-02T13:47:36Z")\
**Posts on this page:** 1\
**Showing post:** 10

<div class="post-metadata">

**Author:** ![bbischof-wandb](https://sea2.discourse-cdn.com/flex020/user_avatar/community.wandb.ai/bbischof-wandb/32/398_2.png) [@bbischof-wandb](https://community.wandb.ai/u/bbischof-wandb)\
**Post date:** [December 5, 2021, 4:50am UTC](https://community.wandb.ai/t/2-jax-course-working-with-neural-networks-in-jax/1455/10 "2021-12-05T04:50:07Z")

</div>

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](https://www.kaggle.com/aakashnain/tf-jax-tutorials-part-9-autodiff-in-jax)

- Here’s the PyTree Notebook from Akash: [TF\_JAX\_tutorials - Part 10 (Pytrees in JAX) | Kaggle](https://www.kaggle.com/aakashnain/tf-jax-tutorials-part-10-pytrees-in-jax)

- Here’s the link to “Flax Basics” from the Flax docs: [Google Colab](https://colab.research.google.com/github/google/flax/blob/main/docs/notebooks/flax_basics.ipynb)

- Here’s the link to “Annotated MNIST” from the Flax docs: [Google Colab](https://colab.research.google.com/github/google/flax/blob/main/docs/notebooks/annotated_mnist.ipynb)

- 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](https://youtu.be/b_6x9BYjmj4)

- Here’s the Scenic implementation of U-net: [scenic/unet.py at main · google-research/scenic · GitHub](https://github.com/google-research/scenic/blob/main/scenic/projects/baselines/unet.py)

---

_[View the full topic](https://community.wandb.ai/t/2-jax-course-working-with-neural-networks-in-jax/1455)._