Notebooks for the "Deep Learning with JAX" book
☆170Jun 7, 2025Updated 8 months ago
Alternatives and similar repositories for Deep_Learning_with_JAX
Users that are interested in Deep_Learning_with_JAX are comparing it to the libraries listed below
Sorting:
- Pointax: PointMaze Environment for JAX☆26Oct 22, 2025Updated 4 months ago
- Single-file SAC-N implementation on jax with flax and equinox. 10x faster than pytorch☆57May 21, 2023Updated 2 years ago
- Jaxpr Visualisation Tool☆36Dec 22, 2024Updated last year
- Official Implementation of `An Optimisation Framework for Unsupervised Environment Design` from RLC 2025☆17Nov 24, 2025Updated 3 months ago
- flexible meta-learning in jax☆16Oct 19, 2023Updated 2 years ago
- Package for data-driven and phenomenological gravitational waveform models☆12Feb 12, 2025Updated last year
- ☆11Dec 6, 2020Updated 5 years ago
- Distrax, but in equinox. Lightweight JAX library of probability distributions and bijectors.☆39Jan 16, 2026Updated last month
- Code related to different aspects of conformal learning☆17Jan 28, 2025Updated last year
- JAX implementation of RL algorithms and vectorized environments☆51Dec 26, 2023Updated 2 years ago
- Detecting compact binary mergers from gravitational wave strain data using neural networks☆18Feb 11, 2026Updated 3 weeks ago
- A Gentle Principled Introduction to Deep Reinforcement Learning☆19Apr 4, 2025Updated 11 months ago
- A process monitoring tool that helps you track and analyze specific processes in real-time.☆16Jan 31, 2025Updated last year
- Flax Implementation of DreamerV3 on Crafter☆18Nov 29, 2025Updated 3 months ago
- ☆566Jul 11, 2024Updated last year
- ⚡ Flashbax: Accelerated Replay Buffers in JAX☆273Sep 22, 2025Updated 5 months ago
- A Large Recurrent Action Model: xLSTM enables Fast Inference for Robotics Tasks☆36Oct 31, 2024Updated last year
- GWInferno: Gravitational-Wave Hierarchical Inference with NumPyro☆21Updated this week
- Models and code for the ICLR 2020 workshop paper "Towards Understanding Normalization in Neural ODEs"☆16Apr 27, 2020Updated 5 years ago
- Performing parameter estimation on gravitational wave data with machine learning☆21Feb 9, 2026Updated 3 weeks ago
- Official pytorch implementation for our ICLR 2023 paper "Latent State Marginalization as a Low-cost Approach for Improving Exploration".☆24Feb 9, 2023Updated 3 years ago
- JAX - A curated list of resources https://github.com/google/jax☆2,058Jan 20, 2026Updated last month
- Hardware-Accelerated Reinforcement Learning Algorithms in pure Jax!☆261Oct 31, 2025Updated 4 months ago
- Add a tqdm progress bar to your JAX scans and loops.☆124May 9, 2025Updated 9 months ago
- A dataloader, but for JAX☆20May 17, 2024Updated last year
- Optax is a gradient processing and optimization library for JAX.☆2,198Feb 26, 2026Updated last week
- Build GPT, Gemma, LlaMa, Mixtral, Whisper, SWin, ViT and more in JAX.☆300Aug 28, 2024Updated last year
- ☆35Nov 22, 2024Updated last year
- An implementation of faster-rcnn for people detection in python☆10Apr 30, 2019Updated 6 years ago
- Repository for the code assignment of the Deep Learning 1 course, Fall 2021 edition☆10Oct 31, 2022Updated 3 years ago
- Rust course☆10May 24, 2025Updated 9 months ago
- JAX-accelerated nuclear equation of state inference and TOV solvers☆10Updated this week
- Source code repository for the AISTAT 2023 paper Transport Reversible Jump Proposals.☆10Mar 3, 2023Updated 3 years ago
- Model Agnostic Meta Learning (MAML) implemented in Flax, the neural network library for JAX.☆21Sep 18, 2020Updated 5 years ago
- Probabilistic Programming and Nested sampling in JAX☆223Updated this week
- Official Implementation of NeurIPS'23 Paper "Cross-Episodic Curriculum for Transformer Agents"☆31Oct 12, 2023Updated 2 years ago
- 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 o…☆776Nov 29, 2023Updated 2 years ago
- make this repo most starred, please 🙏☆11Dec 27, 2024Updated last year
- Ab Initio Energies☆10Nov 22, 2025Updated 3 months ago