Nested Hierarchical Transformer https://arxiv.org/pdf/2105.12723.pdf
created at June 1, 2021, 5:09 a.m.
Elegant easy-to-use neural networks + scientific computing in JAX. https://docs.kidger.site/equinox/
created at July 29, 2021, 2:21 a.m.
JAX/Haiku implementation of "Auction Learning as a Two-Player Game"
created at Aug. 23, 2021, 8:04 a.m.
Unofficial JAX implementations of deep learning research papers
created at Jan. 8, 2022, 1:40 p.m.
The purpose of this repo is to make it easy to get started with JAX, Flax, and Haiku. It contains my "Machine Learning with JAX" series of tutorials (YouTube videos and Jupyter Notebooks) as well as the content I found useful while learning about the JAX ecosystem.
created at Oct. 28, 2021, 8 p.m.
Loopy belief propagation for factor graphs on discrete variables, in JAX!
created at June 9, 2021, 8:21 p.m.
Numerical differential equation solvers in JAX. Autodifferentiable and GPU-capable. https://docs.kidger.site/diffrax/
created at July 25, 2021, 11:42 p.m.
Part 2 of the code for the NeurIPS 2021 paper "Learning Generalized Gumbel-max Causal Mechanisms"
created at Nov. 12, 2021, 6:16 p.m.