JaxPT

GPT models built with JAX

Getting Started

This project implements the GPT series of models using Jax and Flax's NNX library.

Requirements

Install the UV python package managment library

curl -LsSf https://astral.sh/uv/install.sh | sh

Available Commands

The main commands available in the Makefile are:

• make install - Install dependencies from lockfile
• make dev - Install all dependencies including dev from lockfile
• make clean - Clean build artifacts and cache
• make build - Build package
• make lint - Run linting
• make format - Format code
• make lab - Run Jupyter lab server from the project directory

To see all available commands and their descriptions, run: make help