mhlr / awesome-jax
List of awesome JAX resources
☆13Updated 2 years ago
Alternatives and similar repositories for awesome-jax:
Users that are interested in awesome-jax are comparing it to the libraries listed below
- A framework for implementing equivariant DL☆10Updated 3 years ago
- Easily serialize dataclasses to and from tensors (PyTorch, NumPy)☆18Updated 3 years ago
- Simplifying parsing of large jsonline files in NLP Workflows☆12Updated 3 years ago
- JAX implementation of Graph Attention Networks☆13Updated 3 years ago
- ☆33Updated 4 years ago
- LaTeX source code for the slides☆23Updated 3 years ago
- Dive into Jax, Flax, XLA and C++☆31Updated 4 years ago
- ☆30Updated 3 years ago
- Graph Learning with JAX☆14Updated 2 years ago
- 👑 Pytorch code for the Nero optimiser.☆20Updated 2 years ago
- Discontinuous Hamiltonian Monte Carlo in JAX☆41Updated 4 years ago
- An attempt to merge ESBN with Transformers, to endow Transformers with the ability to emergently bind symbols☆15Updated 3 years ago
- Official repository for our ICLR 2021 paper Evaluating the Disentanglement of Deep Generative Models with Manifold Topology☆35Updated 3 years ago
- A runtime shape checker and auto-annotator for tensor programs (pronounced "stanley")☆39Updated 5 years ago
- High-performance tokenized language data-loader for Python C++ extension☆12Updated 6 months ago
- notebooks of cool EBM visualizations☆15Updated 4 years ago
- Quantification of Uncertainty with Adversarial Models☆28Updated last year
- Jax SSM Library☆49Updated 2 years ago
- Stochastic trace estimation using JAX☆12Updated 4 months ago
- Numerically Solving Parametric Families of High-Dimensional Kolmogorov Partial Differential Equations via Deep Learning (NeurIPS 2020)☆22Updated 2 years ago
- Official repository for the paper "Can You Learn an Algorithm? Generalizing from Easy to Hard Problems with Recurrent Networks"☆60Updated 2 years ago
- Toy implementations of some popular ML optimizers using Python/JAX☆44Updated 3 years ago
- ☆12Updated 4 years ago
- A small library for creating and manipulating custom JAX Pytree classes☆57Updated last year
- Fine-grained, dynamic control of neural network topology in JAX.☆21Updated last year
- Official repository for the paper "Neural Differential Equations for Learning to Program Neural Nets Through Continuous Learning Rules" (…☆19Updated 2 years ago
- You should use PySR to find scaling laws. Here's an example.☆33Updated last year
- Companion code for a tutorial on using Hydra.☆29Updated 3 years ago
- Meta-learning inductive biases in the form of useful conserved quantities.☆37Updated 2 years ago
- ☆26Updated 2 years ago