Jit¶

tinax.jit wraps jax.jit with a Chex trace budget validated at wrap time, an optional mesh context, and a jax.vmap fusion for batching. Together bounded_jit and batched_jit are the JAX-native, auto-parallel replacement for legacy pmap-style batching; combine batched_jit with a Manual-axis tinax.parallel.shard_mapped for the manual-SPMD alternative.

from tinax.jit import bounded_jit

train_step = bounded_jit(step, max_traces=1)

static_argnames/donate_argnames overlap or duplication is caught immediately, not on an arbitrary later call.

bounded_jit ¶

bounded_jit(
    function: Callable[P, R],
    *,
    max_traces: int,
    static_argnames: Sequence[str] = (),
    donate_argnames: Sequence[str] = (),
    mesh: Mesh | None = None,
) -> Callable[P, R]

JIT a callable with an exact Chex trace budget, validated at wrap time, not first call.

static_argnames/donate_argnames overlap or duplication is caught here; raw jax.jit only raises on first invocation, which can be arbitrarily deep in a loop. An optional mesh is entered around every call, removing the need for every call site to remember with jax.set_mesh(mesh):.

Parameters:

Name	Type	Description	Default
`function`	`Callable[P, R]`	Callable to JIT-compile.	required
`max_traces`	`int`	Maximum number of distinct traces to permit. Must be non-negative.	required
`static_argnames`	`Sequence[str]`	Keyword argument names to treat as static. Defaults to none.	`()`
`donate_argnames`	`Sequence[str]`	Keyword argument names whose buffers may be donated. Defaults to none.	`()`
`mesh`	`Mesh \| None`	Mesh to enter around every call, or `None` to use the ambient mesh.	`None`

Returns:

Type	Description
`Callable[P, R]`	The trace-budgeted, JIT-compiled callable.

Raises:

Type	Description
`TypeError`	If `function` is not callable, `max_traces` is not an integer (booleans rejected), `static_argnames`/`donate_argnames` is not a sequence of strings, or `mesh` is not a `jax.sharding.Mesh` or `None`.
`ValueError`	If `max_traces` is negative, an argname entry is empty or duplicated, `static_argnames` and `donate_argnames` overlap, or `mesh` is empty.

batched_jit ¶

batched_jit(
    function: Callable[..., R],
    *,
    in_axes: int | None | Sequence[Any] = 0,
    out_axes: Any = 0,
    axis_size: int | None = None,
    max_traces: int,
    mesh: Mesh | None = None,
) -> Callable[..., R]

Batch a callable with jax.vmap, then JIT it with an exact Chex trace budget.

in_axes/out_axes are passed through unvalidated: raw jax.vmap already gives clear, specific errors for a malformed axis specification. axis_size is validated because raw jax.vmap silently accepts a boolean there, producing an unrelated, confusing error instead of a clear rejection. This is the auto-parallel, JAX-native replacement for legacy pmap-style batching; combine with a Manual-axis tinax.parallel.shard_mapped for the manual-SPMD alternative.

Parameters:

Name	Type	Description	Default
`function`	`Callable[..., R]`	Callable to batch and JIT-compile.	required
`in_axes`	`int \| None \| Sequence[Any]`	Input batch axis specification, forwarded to `jax.vmap`. Defaults to 0.	`0`
`out_axes`	`Any`	Output batch axis specification, forwarded to `jax.vmap`. Defaults to 0.	`0`
`axis_size`	`int \| None`	Explicit batch size, or `None` to infer it from the arguments. Must be positive if given.	`None`
`max_traces`	`int`	Maximum number of distinct traces to permit. Must be non-negative.	required
`mesh`	`Mesh \| None`	Mesh to enter around every call, or `None` to use the ambient mesh.	`None`

Returns:

Type	Description
`Callable[..., R]`	The batched, trace-budgeted, JIT-compiled callable.

Raises:

Type	Description
`TypeError`	If `function` is not callable, `axis_size` is not an integer or `None` (booleans rejected), `max_traces` is not an integer (booleans rejected), or `mesh` is not a `jax.sharding.Mesh` or `None`.
`ValueError`	If `axis_size` is not positive, `max_traces` is negative, or `mesh` is empty.