Optax Python Packages

psgd-jax

Implementation of PSGD optimizer in JAX

flax-pilot

A Simplistic trainer for Flax

jqmc

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.

tn4ml

Tensor Networks for Machine Learning

qqn-jax

Quadratic Quasi-Newton (QQN) optimizer, natively implemented in JAX and packaged as a drop-in Optax GradientTransformation.

optax-adan

An implementation of adan optimizer for optax

jaxfss

JAX/Flax implementation of finite-size scaling