optax
Here are 30 public repositories matching this topic...
A library which trains the Fermionic Neural Network to find the ground state wave functions of an atom or a molecule using neural network quantum states.
-
Updated
Jul 18, 2022 - Python
Stochastic Weight Averaging (SWA) transforms for Optax with JAX
-
Updated
Mar 29, 2023 - Python
Simple JAX implementation of Gaussian process regression
-
Updated
Feb 12, 2022
An implementation of MNIST classification using LeNet-300-100 in JAX (using Haiku and Optax).
-
Updated
Aug 7, 2023 - Jupyter Notebook
This repository contains some of the code I wrote for the assignments in DSA4212 - Optimisation for Large-Scale Data-Driven Inference.
-
Updated
May 5, 2023 - Jupyter Notebook
A reimplementation of Parallel DNN Training in JAX by Will Whitney using haiku and optax.
-
Updated
Aug 8, 2023 - Jupyter Notebook
A helper library for training dm-haiku models.
-
Updated
Feb 7, 2021 - Python
Neural implicit digital elevation model
-
Updated
May 25, 2022 - Python
An gradient-free optimization suite written in JAX. We conform to the optax interface and provide ensemble-based optimizers.
-
Updated
Dec 8, 2021 - 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
dm-haiku implementation of hyperbolic neural networks
-
Updated
Jan 29, 2023 - Python
Variational Graph Autoencoder implemented using Jax & Jraph
-
Updated
Mar 22, 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.
-
Updated
Jul 8, 2022 - Python
Goal-conditioned reinforcement learning like 🔥
-
Updated
Feb 3, 2024 - 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."