#
optax
Here are 32 public repositories matching this topic...
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
Mar 24, 2025 - 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
JAX/Flax implementation of finite-size scaling
-
Updated
Jan 15, 2023 - Python
An implementation of adan optimizer for optax
-
Updated
Sep 18, 2022 - 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
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
A Simplistic trainer for Flax
-
Updated
Apr 4, 2025 - 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
An Optax-based JAX implementation of the IVON optimizer for large-scale VI training of NNs (ICML'24 spotlight)
-
Updated
Dec 19, 2024 - Python
dm-haiku implementation of hyperbolic neural networks
-
Updated
Jan 29, 2023 - Python
Direct port of TD3_BC to JAX using Haiku and optax.
-
Updated
Sep 7, 2021 - Python
The (unofficial) vanilla version of WaveRNN
-
Updated
Dec 7, 2021 - 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."