Repository of Jupyter notebook tutorials for teaching the Deep Learning Course at the University of Amsterdam (MSc AI), Fall 2023
-
Updated
Oct 31, 2025 - Jupyter Notebook
#
optax
Here are 47 public repositories matching this topic...
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.
-
Updated
Nov 29, 2023 - Jupyter Notebook
This is the official repository for the paper "Flora: Low-Rank Adapters Are Secretly Gradient Compressors" in ICML 2024.
deep-learning transformers flax lora peft jax random-projection optax large-language-models memory-efficient-tuning
-
Updated
Jul 1, 2024 - Python
Implementation of PSGD optimizer in JAX
-
Updated
Dec 31, 2024 - Python
JAX implementations of various deep reinforcement learning algorithms.
-
Updated
Feb 2, 2025 - Python
Tensor Networks for Machine Learning
python machine-learning python3 flax matrix-product-states tensor-networks jax tensor-network quimb optax
-
Updated
Sep 21, 2025 - Python
JAX/Flax implementation of finite-size scaling
-
Updated
Jan 15, 2023 - Python
Goal-conditioned reinforcement learning like 🔥
goals reinforcement-learning robotics deep-reinforcement-learning flax gymnasium actor-critic deep-deterministic-policy-gradient jax soft-actor-critic optax goal-conditioned-rl distrax gymnasium-robotics mpi4jax droq
-
Updated
Feb 3, 2024 - Python
jQMC code implements two real-space ab initio quantum Monte Carlo (QMC) methods. Variatioinal Monte Carlo (VMC) and lattice regularized diffusion Monte Carlo (LRDMC) methods. jQMC achieves high-performance computations especially on GPUs.
python3 neural-networks gpu-acceleration flax gpu-computing quantum-monte-carlo electronic-structure jax variational-monte-carlo diffusion-monte-carlo ab-initio-simulations optax neural-quantum-states
-
Updated
Dec 5, 2025 - Python
An Optax-based JAX implementation of the IVON optimizer for large-scale VI training of NNs (ICML'24 spotlight)
-
Updated
Dec 19, 2024 - Python
Training methodologies for autoregressive neural operators/emulators in JAX.
deep-learning optimization timeseries-forecasting autoregressive-models optax differentiable-physics unrolled-training
-
Updated
Nov 5, 2024 - Python
An implementation of adan optimizer for optax
-
Updated
Sep 18, 2022 - Python
H-Former is a VAE for generating in-between fonts (or combining fonts). Its encoder uses a Point net and transformer to compute a code vector of glyph. Its decoder is composed of multiple independent decoders which act on a code vector to reconstruct a point cloud representing a glpyh.
fonts neural-network glyphs point-cloud transformer vae flax jax atlasnet point-network font-generation optax
-
Updated
Jul 8, 2022 - Python
JAX implementation of Classical and Quantum Algorithms for Orthogonal Neural Networks by (Kerenidis et al., 2021)
quantum-computing mnist-classification haiku stax jax optax quantum-inspired orthogonal-neural-network
-
Updated
Jan 6, 2023 - Python
Variational Graph Autoencoder implemented using Jax & Jraph
-
Updated
Mar 22, 2024 - Python
A Simplistic trainer for Flax
-
Updated
Apr 4, 2025 - Python
Improve this page
Add a description, image, and links to the optax topic page so that developers can more easily learn about it.
Add this topic to your repo
To associate your repository with the optax topic, visit your repo's landing page and select "manage topics."