optax
Here are 36 public repositories matching this topic...
Simple JAX implementation of Gaussian process regression
-
Updated
Feb 12, 2022
An implementation of adan optimizer for optax
-
Updated
Sep 18, 2022 - Python
sophia optimizer further projected towards flat areas of loss landscape
-
Updated
Aug 19, 2024 - Python
An Optax-based JAX implementation of the IVON optimizer for large-scale VI training of NNs (ICML'24 spotlight)
-
Updated
Jul 30, 2024 - Python
An implementation of MNIST classification using LeNet-300-100 in JAX (using Haiku and Optax).
-
Updated
Aug 7, 2023 - Jupyter Notebook
A helper library for training dm-haiku models.
-
Updated
Feb 7, 2021 - Python
A Simplistic trainer for Flax
-
Updated
Aug 10, 2024 - Python
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
Oxford MSc thesis. PriorVAE with graph convolutional networks for learning locally-aware spatial prior distributions
-
Updated
Apr 16, 2023 - Python
Neural implicit digital elevation model
-
Updated
May 25, 2022 - Python
JAX/Flax implementation of finite-size scaling
-
Updated
Jan 15, 2023 - Python
A place for code to relax in happiness without the pressure of maintenance. Mostly related to Machine Learning and Deep Learning
-
Updated
Aug 9, 2024 - Python
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
Goal-conditioned reinforcement learning like 🔥
-
Updated
Feb 3, 2024 - Python
dm-haiku implementation of hyperbolic neural networks
-
Updated
Jan 29, 2023 - 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
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."