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.array 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¶
| module | Use it for |
|---|---|
array |
NumPy, JAX, DLPack, host copies, logical array inspection, and validated array operations |
jit |
Trace-budgeted compilation, an optional mesh context, and vmap batching |
grad |
Hardened autodiff: value_and_grad, an explicit forward/reverse jacobian, and hessian |
data |
ArrayRecord interchange, dataset splitting, deterministic input pipelines, and multiprocessing worker lifetime |
random |
Typed JAX key validation, coordinate derivation, and key ownership |
debug |
Bounded host observation, completed profiler-call scopes, and sharding visualization |
nn |
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 |
parallel |
Meshes, layouts, shard_map, 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 module in the table above.
- Consult the Public API reference for each module's exact signatures and contracts.