Welcome to the JAX Lab (JAB for short)! This repository contains implementations of foundational Deep Learning architectures in JAX.
| 🚀 Stack: | ||||
|---|---|---|---|---|
Python
|
JAX
|
Haiku
|
Flax
|
Optax
|
- 🐍 Clean and beginner-friendly implementation of Deep Learning algorithms in JAX
- ✍️ Code walkthroughs available here
- Implementing a Transformer Encoder from Scratch with JAX and Haiku 🤖, published in Towards Data Science, Editors' Picks
- Transformer Encoder
- Transformer
To install and set up the project, follow these steps:
-
Clone the repository to your local machine:
git clone https://github.com/RPegoud/jab.git
-
Navigate to the project directory:
cd jab -
Install Poetry (if not already installed):
python -m pip install poetry
-
Install project dependencies using Poetry:
poetry install
-
Activate the virtual environment created by Poetry:
poetry shell