Tinax¶
Reliable boundaries for JAX applications¶
Tinax is a typed library for the JAX ecosystem. It turns high-consequence boundaries into small explicit calls: copying arrays, deriving keys, placing data, managing workers, checkpointing state, and reading weights.
import numpy as np
from tinax.arrays import from_numpy, inspect_array, to_numpy
device = from_numpy(np.arange(8, dtype=np.float32), copy=True)
info = inspect_array(device)
host = to_numpy(device, writable=False)
copy=True is visible. Host materialization is visible. The caller decides when those costs and ownership changes occur.
What Tinax Owns¶
| Domain | Use it for |
|---|---|
arrays |
NumPy, JAX, DLPack, host copies, and logical array inspection |
grain |
Deterministic input pipelines and multiprocessing worker lifetime |
randomness |
Typed JAX key validation, coordinate derivation, and key ownership |
diagnostics |
Bounded host observation and completed profiler-call scopes |
nnx |
Independent Flax NNX graph snapshots, restoration, and copies |
stdlib |
Explicit argparse conversion and isolated stream loggers |
checkpointing |
Atomic Orbax V1 checkpointables and explicit restore targets |
sharding |
Meshes, layouts, placement, process-local data, and NNX integration |
weights |
Tensor manifests and bounded Safetensors interchange |
examples/ contains tested recipes, not stable APIs. Importing tinax alone does not initialize JAX or optional integrations.
Start Here¶
- Read Installation to select the supported Python and JAX environment.
- Read Design Principles for the explicit-policy model and a worked example for every domain in the table above.
- Consult the Public API reference for each domain's exact signatures and contracts.