Tinax

Tinax is a small, typed library of explicit productivity primitives for JAX, Flax NNX, Optax, Orbax, Grain, Chex, and Safetensors workflows.

It provides stable policies for array and RNG ownership, trace-budgeted compilation and batching, hardened autodiff, bounded debug observation, NNX graph copies, explicit stdlib application boundaries, deterministic input pipelines, complete checkpoints, multi-device parallelism, weight interchange, validated metrics, a Weights & Biases run lifecycle, one explicit training loop, and headless explicit-lifecycle plotting. Tested ecosystem recipes live under examples/ without stable API guarantees.

Requirements

• Python 3.12, 3.13, or 3.14

Install

pip install tinax

Install a JAX accelerator distribution when needed. Extra names match JAX's own extras:

pip install "tinax[cuda12]"        # NVIDIA CUDA 12, bundled pip wheels
pip install "tinax[cuda13]"        # NVIDIA CUDA 13, bundled pip wheels
pip install "tinax[cuda12-local]"  # NVIDIA CUDA 12, self-installed CUDA
pip install "tinax[cuda13-local]"  # NVIDIA CUDA 13, self-installed CUDA
pip install "tinax[rocm7-local]"   # AMD ROCm 7
pip install "tinax[tpu]"           # Google Cloud TPU
pip install "tinax[oneapi]"        # Intel GPU through oneAPI

See the installation guide for platform support, tinax[cpu], and accelerator details.

Quick Start

import numpy as np

from tinax.array import from_numpy, inspect_array, to_numpy

host = np.arange(8, dtype=np.float32)
device = from_numpy(host, copy=True)
info = inspect_array(device)
round_trip = to_numpy(device, writable=False)

Importing tinax alone is inert. Import the module that owns the behavior you need.

Documentation

• Documentation
• Changelog
• Contributing
• Security policy
• Release guide

License

Apache-2.0. See LICENSE.