Repository of Jupyter notebook tutorials for teaching the Deep Learning Course at the University of Amsterdam (MSc AI), Fall 2023
-
Updated
Jun 1, 2026 - Jupyter Notebook
#
optax
Here are 63 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
Jun 23, 2026 - 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
Jul 7, 2026 - Python
JAX/Flax implementation of finite-size scaling
-
Updated
Jan 15, 2023 - Python
Training methodologies for autoregressive neural operators/emulators in JAX.
deep-learning optimization timeseries-forecasting autoregressive-models optax differentiable-physics unrolled-training
-
Updated
Jun 1, 2026 - 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
EEG task classification with CNN, LSTM, CNN-LSTM, and GAN augmentation across TensorFlow, PyTorch, and JAX.
-
Updated
May 23, 2026 - 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
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
Training Quantum Signal Processing phase angles from scratch using JAX gradient descent and PennyLane.
python demo qml quantum jupyter-notebook quantum-computing qsp gradient-descent quantum-algorithms quantum-machine-learning jax chebyshev-approximation pennylane optax quantum-signal-processing qsvt variational-quantum
-
Updated
Jun 12, 2026 - Jupyter Notebook
dLLM training implementation on pure jax/flax (w/o pytorch) for Google TPUs(v4/v5e/v6e). #TPUSprint #TRC
-
Updated
Apr 25, 2026 - 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
An implementation of adan optimizer for optax
-
Updated
Sep 18, 2022 - Python
Add this topic to your repo
To associate your repository with the optax topic, visit your repo's landing page and select "manage topics."