Projects

latest | popular

Getting started with JAX (MLPs, CNNs & RNNs)
Learn the building blocks of JAX and use them to build some standard Deep Learning architectures (MLP, CNN, RNN, etc.).
jax xla autograd tpu
Finetuning Transformers with JAX + Haiku
Walking through a port of the RoBERTa pre-trained model to JAX + Haiku, then fine-tuning the model to solve a downstream task.
jax haiku roberta transformers
Flax: Google’s Open Source Approach To Flexibility In ML
A gentle introduction to Flax: a neural network library for JAX that is designed for flexibility.
flax jax deep-learning library
SymJAX
A symbolic CPU/GPU/TPU programming
jax xla autograd symjax
Convoluted Stuff
Optimising compilers, and how thousand-year-old math shaped deep learning.
matrix-multiplication convolutional-neural-networks jax im2col
Implementing Graph Neural Networks with JAX
I’ll talk about my experience on how to build and train Graph Neural Networks (GNNs) with JAX.
graph-neural-networks jax graphs tutorial
From PyTorch to JAX
Towards neural net frameworks that purify stateful code.
jax haiku tutorial
Lagrangian Neural Networks
Trying to learn a simulation? Try Lagrangian Neural Networks, which explicitly conserve energy and may generalize better!
deep-learning graph-neural-networks jax interpretability
Using JAX to Improve Separable Image Filters
Optimizing the filters to improve the filtered images for computer vision tasks.
jax numpy computer-vision separable-filters
projects 1 - 9 of 9
Topic experts
Share a project
Share something interesting you found that's made with ML.