Jupyter Notebook Jax

Open-source Jupyter Notebook projects categorized as Jax
Edit details
flax Deep Learning Tensorflow Machine Learning Pytorch
SaaSHub - Software Alternatives and Reviews
SaaSHub helps you find the best software and product alternatives
www.saashub.com
featured

Top 20 Jupyter Notebook Jax Projects

Jupyter Notebook logo
Jax

  1. flax

    Flax is a neural network library for JAX that is designed for flexibility.

    Project mention: Flax: A neural network library and ecosystem for Jax designed for flexibility | news.ycombinator.com | 2025-12-15

  2. SaaSHub

    SaaSHub - Software Alternatives and Reviews. SaaSHub helps you find the best software and product alternatives

    SaaSHub logo

  3. pyprobml

    Python code for "Probabilistic Machine learning" book by Kevin Murphy

  4. whisper-jax

    JAX implementation of OpenAI's Whisper model for up to 70x speed-up on TPU.

  5. brax

    Massively parallel rigidbody physics simulation on accelerator hardware.

  6. uvadlc_notebooks

    Repository of Jupyter notebook tutorials for teaching the Deep Learning Course at the University of Amsterdam (MSc AI), Fall 2023

  7. score_sde

    Official code for Score-Based Generative Modeling through Stochastic Differential Equations (ICLR 2021, Oral)

  8. jax-cfd

    Computational Fluid Dynamics in JAX

  9. tensor-sensor

    The goal of this library is to generate more helpful exception messages for matrix algebra expressions for numpy, pytorch, jax, tensorflow, keras, fastai.

  10. get-started-with-JAX

    The purpose of this repo is to make it easy to get started with JAX, Flax, and Haiku. It contains my "Machine Learning with JAX" series of tutorials (YouTube videos and Jupyter Notebooks) as well as the content I found useful while learning about the JAX ecosystem.

  11. jaxrl

    JAX (Flax) implementation of algorithms for Deep Reinforcement Learning with continuous action spaces.

  12. felafax

    Felafax is building AI infra for non-NVIDIA GPUs

  13. TF_JAX_tutorials

    All about the fundamental blocks of TF and JAX!

  14. indaba-pracs-2022

    Notebooks for the Practicals at the Deep Learning Indaba 2022.

  15. hyper-nn

    Easy Hypernetworks in Pytorch and Jax

  16. jaxrenderer

    Differentiable Rasteriser implemented in JAX. Reference: https://github.com/erwincoumans/tinyrenderer, https://github.com/ssloy/tinyrenderer/wiki; PR: https://github.com/google/brax/pull/367

  17. astronomix

    differentiable (magneto)hydrodynamics for astrophysics in JAX

    Project mention: Run Your Own Astrophysical Fluid Sims in Colab | news.ycombinator.com | 2025-12-02

  18. diajax

    Dia-JAX: A JAX port of Dia, the text-to-speech model for generating realistic dialogue from text with emotion and tone control.

  19. forwardalgorithms

    A tensorflow implementation of the Forward-Forward Algorithm from NeurIPS '22.

  20. spinq-vqe

    VQE simulation of the Kagome antiferromagnet (Mn₃Sn) — HEA & MERA ansatze, entanglement entropy, ED benchmarks, SOC QAOA.

    Project mention: Simulating Antiferromagnets with PennyLane: A Deep Dive into spinq-vqe | dev.to | 2026-06-25

    Quantum Machine Learning (QML) is rapidly moving from theoretical whitepapers into practical engineering pipelines. If you are curious about how quantum algorithms actually interface with material science, the newly open-sourced spinq-vqe repository from ARPA's QONDRA division is a perfect sandbox.

  21. trainable-qsp-angles

    Reproducible QSP benchmark: gradient-learned phase angles on a JAX flat circuit vs mapped analytic solvers.

    Project mention: You Can Now Train Quantum Signal Processing Phase Angles with Gradient Descent | dev.to | 2026-06-11

    A new paper and open‑source implementation (DOI: 10.5281/zenodo.20645403, github.com/rosspeili/trainable-qsp-angles) shows that you can learn QSP phase angles from random initialization using gradient descent – with no analytic pre‑computation.

NOTE: The open source projects on this list are ordered by number of github stars. The number of mentions indicates repo mentiontions in the last 12 Months or since we started tracking (Dec 2020).

Jupyter Notebook Jax discussion

Log in or Post with

Jupyter Notebook Jax related posts

Index

What are some of the best open-source Jax projects in Jupyter Notebook? This list will help you:

# Project Stars
1 flax 7,261
2 pyprobml 7,106
3 whisper-jax 4,684
4 brax 3,202
5 uvadlc_notebooks 3,167
6 score_sde 1,828
7 jax-cfd 950
8 tensor-sensor 805
9 get-started-with-JAX 783
10 jaxrl 756
11 felafax 570
12 TF_JAX_tutorials 280
13 indaba-pracs-2022 178
14 hyper-nn 106
15 jaxrenderer 79
16 astronomix 62
17 diajax 29
18 forwardalgorithms 10
19 spinq-vqe 4
20 trainable-qsp-angles 2

Sponsored
SaaSHub - Software Alternatives and Reviews
SaaSHub helps you find the best software and product alternatives
www.saashub.com

Did you know that Jupyter Notebook is
the 15th most popular programming language
based on number of references?

Loading...