Training GPT-2 on FineWeb-Edu in JAX/Flax

A FLAX NNX implementation of GPT2

Flax Engine Documentation compiled for static hosting

Machine learning samples

Graphs the activity of a flax harvester in a linux terminal.

Unfinished library for the NLP tasks using Jax and flax

LSTM POS Tagger implementation in Jax and Flax

This repository implements sentiment classifier with Google Jax using Bert transformer as backbone. It also shows model checkpointing and loading for inference.

An implementation of Legendre Memory Units in Flax/Jax.

so-vits-svc rewritten in jax and flax.

SQA scottish Science Baccalaureate of Christopher Rae. Improving the preformance of distributed data parallelism of low wifi bandwidths

Common practices for distributed training using various backends

Yet Another X: JAX/FLAX module tracing, modification, and evaluation.

Whisper Transcription Service

GPT from scratch (tensorflow / pytorch / jax)

This repository contains some of the code I wrote for the assignments in DSA4212 - Optimisation for Large-Scale Data-Driven Inference.