Parallel¶

tinax.parallel is the home for JAX's multi-device story: meshes, explicit layouts, host placement, process-local arrays, Manual-axis shard_map, and Flax NNX sharding integration.

Process-local construction is not a gather. Callers must ensure replica consistency when their layout requires it. layout builds concrete Explicit-axis layouts; shard_mapped is the Manual-axis counterpart for callables that issue explicit collectives, validated against the mesh before staging.

See the Distributed Placement guide for task-oriented usage.

make_mesh ¶

make_mesh(
    devices: Sequence[Device],
    shape: Sequence[int],
    axis_names: Sequence[str],
    *,
    axis_types: Sequence[AxisType],
) -> Mesh

Build a mesh from exactly the caller's devices and explicit axis policy.

Parameters:

Name	Type	Description	Default
`devices`	`Sequence[Device]`	Non-empty sequence of distinct `jax.Device` instances, used exactly once each.	required
`shape`	`Sequence[int]`	Positive size per mesh axis. Its product must equal the device count.	required
`axis_names`	`Sequence[str]`	Unique, non-empty name per mesh axis.	required
`axis_types`	`Sequence[AxisType]`	`jax.sharding.AxisType` per mesh axis.	required

Returns:

Type	Description
`Mesh`	A `jax.sharding.Mesh` over the provided devices.

Raises:

Type	Description
`TypeError`	If an argument is not the expected sequence or element type (booleans rejected for sizes).
`ValueError`	If devices are empty or duplicated, a size is not positive, names are empty or not unique, the lengths disagree, or the shape does not use every device exactly once.

layout ¶

layout(
    mesh: Mesh,
    partitions: Sequence[str | tuple[str, ...] | None],
) -> NamedSharding

Create an Explicit named layout whose rank is inferred from its partition entries.

Parameters:

Name	Type	Description	Default
`mesh`	`Mesh`	Mesh with Explicit axes to build the layout over.	required
`partitions`	`Sequence[str \| tuple[str, ...] \| None]`	One entry per array dimension: a mesh-axis name, a non-empty tuple of names, or `None` for replication.	required

Returns:

Type	Description
`NamedSharding`	A `jax.sharding.NamedSharding` for the described layout.

Raises:

Type	Description
`TypeError`	If `mesh` is not a `jax.sharding.Mesh` or an entry has the wrong type.
`ValueError`	If `mesh` does not use Explicit axes, a partition axis is empty or unknown, or a mesh axis partitions more than one dimension.

shard_mapped ¶

shard_mapped(
    function: Callable[P, R],
    *,
    mesh: Mesh,
    in_specs: object,
    out_specs: object,
    axis_names: frozenset[str] = frozenset(),
    check_vma: bool = True,
) -> Callable[P, R]

Manually shard-map a callable, validating specs against Manual mesh axes before staging.

layout builds concrete Explicit-axis layouts; shard_mapped is the Manual-axis counterpart for callables that issue explicit collectives. Every mesh axis this call is manual over (axis_names, or every mesh axis if empty, matching jax.shard_map's own semantics) must be AxisType.Manual; every in_specs/out_specs leaf may only reference axes within that set. Validation runs at wrap time, before any argument is staged.

Parameters:

Name	Type	Description	Default
`function`	`Callable[P, R]`	Callable to shard-map. Each call receives one shard of its mapped arguments.	required
`mesh`	`Mesh`	Mesh whose axes this call is manual over.	required
`in_specs`	`object`	Pytree of `jax.sharding.PartitionSpec` leaves describing argument sharding.	required
`out_specs`	`object`	Pytree of `jax.sharding.PartitionSpec` leaves describing output sharding.	required
`axis_names`	`frozenset[str]`	Mesh axes this call is manual over. Empty means every mesh axis.	`frozenset()`
`check_vma`	`bool`	Whether to enable `jax.shard_map`'s replication validity checks. Defaults to `True`.	`True`

Returns:

Type	Description
`Callable[P, R]`	The shard-mapped callable.

Raises:

Type	Description
`TypeError`	If `function` is not callable, `mesh` is not a `jax.sharding.Mesh`, `axis_names` is not a `frozenset` of strings, `check_vma` is not a bool, or an `in_specs`/`out_specs` leaf is not a `jax.sharding.PartitionSpec`.
`ValueError`	If `mesh` is empty, `axis_names` is not a subset of the mesh's axes, an effective manual axis is not `AxisType.Manual`, or a spec references an axis outside the effective manual set.

place_host_array ¶

place_host_array(
    array: ndarray, sharding: NamedSharding, *, copy: bool
) -> Array

Place host data asynchronously; copy snapshots now, while false permits a lifetime source alias.

Parameters:

Name	Type	Description	Default
`array`	`ndarray`	Host NumPy array to place.	required
`sharding`	`NamedSharding`	Fully addressable `jax.sharding.NamedSharding` describing placement.	required
`copy`	`bool`	If `True`, snapshot `array` before the asynchronous transfer. If `False`, permit aliasing the source for the transfer's lifetime.	required

Returns:

Type	Description
`Array`	A JAX array placed according to `sharding`.

Raises:

Type	Description
`TypeError`	If `array` is not a host NumPy array, `sharding` is not a `NamedSharding`, or `copy` is not a bool.
`ValueError`	If `sharding` is not fully addressable or cannot partition `array`'s shape.

from_process_local_data ¶

from_process_local_data(
    local_data: ndarray,
    sharding: NamedSharding,
    global_shape: tuple[int, ...],
    *,
    copy: bool,
    replica_policy: Literal[
        "disallow", "assume_consistent"
    ],
    process_index: int | None = None,
    process_count: int | None = None,
) -> Array

Create without gathering; copy controls source lifetime and assume_consistent attests replica equality.

Constructs a global array from each process's local shard. This is not a gather: callers are responsible for replica consistency.

Parameters:

Name	Type	Description	Default
`local_data`	`ndarray`	This process's local host shard.	required
`sharding`	`NamedSharding`	Target `jax.sharding.NamedSharding`.	required
`global_shape`	`tuple[int, ...]`	Logical shape of the assembled global array.	required
`copy`	`bool`	If `True`, snapshot `local_data`; if `False`, permit aliasing it.	required
`replica_policy`	`Literal['disallow', 'assume_consistent']`	`"disallow"` rejects cross-process replicas; `"assume_consistent"` attests that replicated data is identical.	required
`process_index`	`int \| None`	Optional declared process index; must match the JAX runtime.	`None`
`process_count`	`int \| None`	Optional declared process count; must match the JAX runtime.	`None`

Returns:

Type	Description
`Array`	The assembled global JAX array.

Raises:

Type	Description
`TypeError`	If an argument has the wrong type (booleans rejected for integers).
`ValueError`	If ranks disagree, `replica_policy` is invalid, an index or count is out of range, the sharding cannot partition `global_shape`, local extents are inconsistent, or cross-process replicas exist under `"disallow"`.
`RuntimeError`	If a declared `process_index` or `process_count` does not match the JAX runtime.

logical_payload_nbytes ¶

logical_payload_nbytes(array: Array) -> int | None

Return nominal logical bytes, or None when the dtype has no standard byte representation.

Parameters:

Name	Type	Description	Default
`array`	`Array`	JAX array to size. Its data is not materialized.	required

Returns:

Type	Description
`int \| None`	The logical payload size in bytes, or `None` for extended dtypes.

Raises:

Type	Description
`TypeError`	If `array` is not a `jax.Array`.

addressable_payload_nbytes ¶

addressable_payload_nbytes(array: Array) -> int | None

Return nominal addressable bytes including replicas, or None for an extended dtype.

Parameters:

Name	Type	Description	Default
`array`	`Array`	JAX array to size. Its data is not materialized.	required

Returns:

Type	Description
`int \| None`	The total bytes across this process's addressable shards (replicas counted), or
`int \| None`	`None` for extended dtypes.

Raises:

Type	Description
`TypeError`	If `array` is not a `jax.Array`.

addressable_indices ¶

addressable_indices(
    array: Array,
) -> dict[Device, tuple[slice | int, ...] | None]

Return each addressable device's global logical index.

Parameters:

Name	Type	Description	Default
`array`	`Array`	JAX array whose sharding is inspected.	required

Returns:

Type	Description
`dict[Device, tuple[slice \| int, ...] \| None]`	A mapping from each addressable `jax.Device` to the global index tuple of the
`dict[Device, tuple[slice \| int, ...] \| None]`	shard it holds, or `None` where unindexed.

Raises:

Type	Description
`TypeError`	If `array` is not a `jax.Array`.

abstract_sharded_state ¶

abstract_sharded_state(
    initializer: Callable[[], T], mesh: Mesh
) -> GraphState[T]

Create sharding-aware abstract graph state without allocating parameters.

Parameters:

Name	Type	Description	Default
`initializer`	`Callable[[], T]`	Zero-argument callable that constructs NNX state.	required
`mesh`	`Mesh`	Homogeneous mesh (all Auto or all Explicit axes, no Manual).	required

Returns:

Type	Description
`GraphState[T]`	A `flax.nnx.GraphState` describing the abstract, sharding-aware state.

Raises:

Type	Description
`TypeError`	If `initializer` is not callable, `mesh` is not a `jax.sharding.Mesh`, or the initializer does not return an NNX graph node.
`ValueError`	If `mesh` is empty, uses Manual axes, or mixes Auto and Explicit axes.

eager_sharded_init ¶

eager_sharded_init(
    initializer: Callable[[], T], mesh: Mesh
) -> T

Run an NNX initializer once under a homogeneous mesh with eager sharding enabled.

Parameters:

Name	Type	Description	Default
`initializer`	`Callable[[], T]`	Zero-argument callable that constructs and returns NNX state.	required
`mesh`	`Mesh`	Homogeneous mesh (all Auto or all Explicit axes, no Manual).	required

Returns:

Type	Description
`T`	The initializer's result, sharded according to `mesh`.

Raises:

Type	Description
`TypeError`	If `initializer` is not callable or `mesh` is not a `jax.sharding.Mesh`.
`ValueError`	If `mesh` is empty, uses Manual axes, or mixes Auto and Explicit axes.

state_shardings ¶

state_shardings(
    state: StateMapping, mesh: Mesh
) -> StateMapping

Derive named shardings from one model or optimizer state independently.

Parameters:

Name	Type	Description	Default
`state`	`StateMapping`	Mapping of NNX state (for example, model or optimizer state).	required
`mesh`	`Mesh`	Homogeneous mesh (all Auto or all Explicit axes, no Manual).	required

Returns:

Type	Description
`StateMapping`	A mapping of the same structure whose leaves are the derived named shardings.

Raises:

Type	Description
`TypeError`	If `state` is not a mapping or `mesh` is not a `jax.sharding.Mesh`.
`ValueError`	If `mesh` is empty, uses Manual axes, or mixes Auto and Explicit axes.