Taktiny

A deep learning library built on JAX.

Define models with Python modules, transform them with JAX, and choose how to load data, distribute computation, and run training.

Installation · Quickstart · API reference

Taktiny is experimental. APIs may change as the library develops.

A model is a PyTree

Layers compose as ordinary Python objects. A module can be passed to a JAX transformation alongside its inputs.

import jax
import jax.numpy as jnp
from taktiny import nn


class MLP(nn.Module):
    def __init__(self, *, rngs: nn.Rngs):
        self.hidden = nn.Linear(16, 32, rngs=rngs)
        self.output = nn.Linear(32, 4, rngs=rngs)

    def __call__(self, x):
        return self.output(jax.nn.relu(self.hidden(x)))


model = MLP(rngs=nn.Rngs(0))
forward = jax.jit(model)

y = forward(jnp.ones((8, 16)))
print(y.shape)  # (8, 4)

The quickstart continues with data loading, optimization, and evaluation.

Find your way

Start here

Set up your environment and train a first model.

• Install Taktiny

• Build and train a model

Guides

Work with the parts of a training pipeline.

• Saving and Loading Models

• Parameter-efficient fine-tuning

• Data loading and preprocessing

• Training and checkpoints

• Sharding and parallelism

Tutorials

Complete examples, from simple models to image generation.

• Linear regression

• Image classification

• Generative adversarial networks

API reference

Signatures, arguments, and examples by package.

• Neural network modules

• Adapter injection

• Data transforms and loaders

• Trainer and callbacks

• JAX and quantization utilities