optax

Star

Here are 41 public repositories matching this topic...

Language:All

Filter by language

All41 Python31 Jupyter Notebook8 TeX1

Sort:Most stars

Sort options

Most stars Fewest stars Most forks Fewest forks Recently updated Least recently updated

phlippe /uvadlc_notebooks

Star2.8k

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

tutorial deep-learning tutorials pytorch flax jax pytorch-lightning optax uvadlc

UpdatedMar 16, 2025
Jupyter Notebook

gordicaleksa /get-started-with-JAX

Star704

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.

python machine-learning tutorial deep-learning jupyter numpy flax haiku lax jax xla optax learn-jax

UpdatedNov 29, 2023
Jupyter Notebook

erfanzar /EasyDeL

Star251

Accelerate, Optimize performance with streamlined training and serving options with JAX.

transformers flax gpt jax optax easydel

UpdatedMar 18, 2025
Python

BorealisAI /flora-opt

Star102

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

UpdatedJul 1, 2024
Python

varun-ml /diffusion-models-tutorial

Star57

Experiment with diffusion models that you can run on your local jupyter instances

pytorch vae jax diffusion-models dm-haiku optax

UpdatedOct 27, 2024
Jupyter Notebook

evanatyourservice /psgd_jax

Star29

Implementation of PSGD optimizer in JAX

optimization preconditioning jax second-order-optimization optax psgd

UpdatedDec 31, 2024
Python

hushon /JAX-ResNet-CIFAR10

Star23

Simple CIFAR10 ResNet example with JAX.

deep-learning resnet cifar10 haiku jax optax

UpdatedJun 1, 2021
Python

hamishs /JAX-RL

Star21

JAX implementations of various deep reinforcement learning algorithms.

reinforcement-learning haiku jax optax

UpdatedFeb 2, 2025
Python

bsc-quantic /tn4ml

Star17

Tensor Networks for Machine Learning

machine-learning flax jax quimb optax tensornetworks

UpdatedMar 14, 2025
Python

frankroeder /goal_conditioned_rl

Star11

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

UpdatedFeb 3, 2024
Python

JesseFarebro /flax-mup

Star10

Maximal Update Parametrization (μP) with Flax & Optax.

flax jax optax

UpdatedDec 27, 2023
Python

yonesuke /jaxfss

Star9

JAX/Flax implementation of finite-size scaling

machine-learning deep-learning neural-network physics flax jax optax

UpdatedJan 15, 2023
Python

hr0nix /optax-adan

Star7

An implementation of adan optimizer for optax

deeplearning optimization-methods optimization-algorithms jax optax

UpdatedSep 18, 2022
Python

Ceyron /trainax

Star7

Training methodologies for autoregressive neural operators/emulators in JAX.

deep-learning optimization timeseries-forecasting autoregressive-models optax differentiable-physics unrolled-training

UpdatedNov 5, 2024
Python

NITHISHM2410 /flax-pilot

Star7

A Simplistic trainer for Flax

machine-learning deep-learning flax jax tensorflow2 optax flax-linen

UpdatedFeb 11, 2025
Python

mzguntalan /h-former

Star7

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