tinax: primitives for the JAX ecosystem

tinax is a library of small, typed functions for the seams between JAX, Flax NNX, Optax, Orbax, Grain, Chex, and Safetensors. Each seam hides a policy; each tinax function makes that policy a required argument. Importing tinax on its own does nothing. You import the module that owns the behaviour you need, and CI installs the built wheel into a fresh environment to prove that import tinax never pulls in JAX.

What it covers

Twelve modules, each a thin public interface over private implementations:

• array: from_numpy(copy=...) and to_numpy(writable=...) make the caller choose aliasing or copying; safe_astype rejects lossy casts unless told otherwise.
• jit: bounded_jit(max_traces=N) wraps jax.jit with a trace budget and checks static_argnames against donate_argnames at wrap time, not on the first call.
• grad: jacobian requires an explicit mode="forward" or "reverse".
• random: derive_process_step_key fixes the fold-in order (process, step, stream) and bounds each coordinate to 32 bits.
• checkpointing: a wrapper over Orbax with immutable destinations.
• parallel: mesh construction separated from placement, with shard_map axis types validated.
• weights: load_safetensors reads the manifest first and enforces a byte budget.

Errors follow one contract: TypeError for the wrong category, ValueError for an invalid value, and a bool is never accepted where an int is expected.

Discipline

The design document says what the library is not: a framework around JAX. One module, tinax.lax, was considered and rejected because JAX already raises clear errors there. The rejection is written down so it does not come back.

Dependencies are pinned exactly and locked. CI runs the tests, ruff, and the ty type checker on three operating systems and three Python versions, builds the wheel, checks it with twine, builds the docs in strict mode, and only then releases. Seventeen unit test files and nine integration files cover checkpoint resume, data workers, DLPack interchange, safetensors, shard_map, and NNX parallelism; accelerator and multi-host tests have their own pytest markers. Twelve example recipes ship under examples/ and are tested, with no stability promise.