google / google/flax

Add more examples in the NNX docs

Open
#5,433 4 comments 0 reactions 2 assignees Claimed by @samanklesaria View on GitHub
Dominant language
Jupyter Notebook
Stars
7.3k
Forks
833
Avg merge
5h 11m
Merged PRs (30d)
5

Description

We plan for moving Jax AI Stack examples to NNX docs:

- [ ] (Low prio) Part 2: Debug a variational autoencoder (VAE)
- [x] Part 3: Train a diffusion model for image generation (@samanklesaria, https://github.com/google/flax/pull/5403)
- [x] Visualize JAX model metrics with TensorBoard (@samanklesaria, https://github.com/google/flax/pull/5425)
- [x] Introduction to Data Loaders on CPU with JAX / Introduction to Data Loaders on GPU with JAX (@samanklesaria #5454 )
- [x] JAX for PyTorch users / Porting a PyTorch model to JAX (@vfdev-5, https://github.com/google/flax/pull/5408)
- [x] Train a miniGPT language model with JAX (@samanklesaria, https://github.com/google/flax/pull/5405)
- [ ] Text classification with a transformer language model using JAX
- [x] Machine Translation with encoder-decoder transformer model (@samanklesaria, https://github.com/google/flax/pull/5431)
- [x] Image segmentation with UNETR model (@vfdev-5, https://github.com/google/flax/pull/5463)
- [ ] ~Image Captioning with Vision Transformer (ViT) model (@vfdev-5, ...)~
- [x] Object Detection with DETR, inference (@vfdev-5 )
- [ ] training more recent version of DETR, e.g. RF-DETR (but more code to write, can be very complex notebook)
- [x] Train a Vision Transformer (ViT) for image classification with JAX (@vfdev-5, https://github.com/google/flax/pull/5455)
- [ ] Time series classification with CNN

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.