6 projects
dm-clrs
The CLRS Algorithmic Reasoning Benchmark.
optax
A gradient processing and optimisation library in JAX.
dm-pix
PIX is an image processing library in JAX, for JAX.
chex
Chex: Testing made fun, in JAX!
distrax
Distrax: Probability distributions in JAX.
rlax
A library of reinforcement learning building blocks in JAX.