Why tinax does nothing when you import it

import tinax does nothing. It does not import JAX, it does not register anything, it does not configure a backend. The CI pipeline installs the built wheel into three fresh virtual environments and asserts exactly that. If a future change made the top-level import pull in JAX, the release would fail before it shipped.

That rule is the whole library in miniature. tinax is narrow on purpose. It is not a framework around JAX. It makes the policies at the seams between JAX, Flax NNX, Optax, Orbax, Grain, Chex, and Safetensors explicit, and it stops there.

Boundaries have policies

Each of those libraries is well designed on its own. The trouble is at the seams, where a decision gets made and nobody says so.

When you convert a NumPy array to a JAX array, is the memory shared or copied? The answer depends on the backend and the dtype, and the default is silent. tinax's from_numpy takes copy=True or copy=False and will not run without one. to_numpy asks the same question from the other side with writable.

When you cast an array, can the cast lose information? safe_astype refuses a lossy cast unless you pass allow_lossy=True. The check is cheap. The alternative is finding a silent float64-to-float16 truncation three weeks later.

When you jit a function, how many times is it allowed to retrace before something is wrong? bounded_jit(max_traces=N) gives the answer a number. It also checks that static_argnames and donate_argnames do not overlap when the function is wrapped, not on the first call, because the first call is often inside a training loop on a machine you are paying for by the hour.

from tinax.array import from_numpy
from tinax.jit import bounded_jit
 
device = from_numpy(host, copy=True)          # aliasing is a decision, not a default
step = bounded_jit(train_step, max_traces=2)  # retrace more than twice and it raises

The same shape repeats through the library. jacobian will not guess between forward and reverse mode. derive_process_step_key fixes the order in which process, step, and stream are folded into a PRNG key: $k = \mathrm{fold}(\mathrm{fold}(\mathrm{fold}(k_0, p), s), t)$, with each of $p$, $s$, $t$ bounded to $2^{32}$. Two people who choose different orders get different random streams, and neither of them is wrong. load_safetensors reads the manifest before the tensors and enforces a byte budget. Checkpoints go to immutable destinations. Mesh construction is one step and placement is another.

Errors are part of the interface

Wrong category of value raises TypeError. Right category, invalid value raises ValueError. A Python bool is never accepted where an int is expected, because True is a valid integer in Python and that has caused enough bugs already.

Knowing what not to add

There was going to be a tinax.lax module. It was evaluated for the 0.1.3 release and rejected: JAX already raises clear errors there, and a wrapper would have added a layer without adding a decision. The rejection is written down in the design document so the idea does not come back every few months.

Proving it

The test suite is seventeen unit files and nine integration files: checkpoint resume, data workers, DLPack interchange, safetensors, shard_map, and NNX parallelism, with accelerator and multi-host tests behind their own markers. CI runs pytest, ruff, and the ty type checker across three operating systems and three Python versions, builds the wheel, checks it with twine, builds the documentation in strict mode, and then does the import check that this post opened with.

Dependencies are pinned exactly and locked. That is unusual for a library, and it is on purpose: the seams tinax covers move between JAX releases, and a version that claims to work with "any recent JAX" would be lying.