google-deepmind / eigengameLinks
Open source code for EigenGame.
☆30Updated 2 years ago
Alternatives and similar repositories for eigengame
Users that are interested in eigengame are comparing it to the libraries listed below
Sorting:
- Parameter-Free Optimizers for Pytorch☆130Updated last year
- ☆114Updated 3 weeks ago
- Lightning-like training API for JAX with Flax☆41Updated 6 months ago
- Multi-framework implementation of Deep Kernel Shaping and Tailored Activation Transformations, which are methods that modify neural netwo…☆70Updated 3 weeks ago
- Turn jitted jax functions back into python source code☆22Updated 6 months ago
- Riemannian Optimization Using JAX☆49Updated last year
- This is a port of Mistral-7B model in JAX☆32Updated 11 months ago
- Implementations and checkpoints for ResNet, Wide ResNet, ResNeXt, ResNet-D, and ResNeSt in JAX (Flax).☆112Updated 3 years ago
- Euclidean Wasserstein-2 optimal transportation☆47Updated last year
- Meta-learning inductive biases in the form of useful conserved quantities.☆37Updated 2 years ago
- JAX implementation of Learning to learn by gradient descent by gradient descent☆27Updated 8 months ago
- The Energy Transformer block, in JAX☆58Updated last year
- Perceiver (transformer variant) implemented in JAX and Flax☆12Updated 4 years ago
- Laplace Redux -- Effortless Bayesian Deep Learning☆42Updated 2 weeks ago
- ☆43Updated 3 weeks ago
- Code accompanying our paper "Feature Learning in Infinite-Width Neural Networks" (https://arxiv.org/abs/2011.14522)☆62Updated 4 years ago
- Pytorch-like dataloaders for JAX.☆87Updated last month
- A selection of neural network models ported from torchvision for JAX & Flax.☆44Updated 4 years ago
- Laplace approximated BNN surrogate for BoTorch☆8Updated 7 months ago
- Pytorch implementation of preconditioned stochastic gradient descent (Kron and affine preconditioner, low-rank approximation precondition…☆177Updated 2 weeks ago
- Agustinus' very opiniated publication-ready plotting library☆66Updated last month
- This repository contains a Jax implementation of conformal training corresponding to the ICLR'22 paper "learning optimal conformal classi…☆130Updated 2 years ago
- code for "Neural Conservation Laws A Divergence-Free Perspective".☆38Updated 2 years ago
- ☆19Updated 3 years ago
- Unofficial but Efficient Implementation of "Mamba: Linear-Time Sequence Modeling with Selective State Spaces" in JAX☆83Updated last year
- Implementation of PSGD optimizer in JAX☆33Updated 5 months ago
- Neural Networks for JAX☆84Updated 9 months ago
- [TMLR 2022] Curvature access through the generalized Gauss-Newton's low-rank structure: Eigenvalues, eigenvectors, directional derivative…☆17Updated last year
- Pytorch code for "Improving Self-Supervised Learning by Characterizing Idealized Representations"☆41Updated 2 years ago
- Simple (and cheap!) neural network uncertainty estimation☆66Updated last month