jax.sharding module

Classes

class jax.sharding.Sharding(*args, **kwargs)

Describes how a jax.Array is laid out across devices.

property addressable_devices: set[Device]

The set of devices in the Sharding that are addressable by the current process.

addressable_devices_indices_map(global_shape)

A mapping from addressable devices to the slice of array data each contains.

addressable_devices_indices_map contains that part of device_indices_map that applies to the addressable devices.

Parameters:

global_shape (Shape)

Return type:

Mapping[Device, Index | None]

property device_set: set[Device]

The set of devices that this Sharding spans.

In multi-controller JAX, the set of devices is global, i.e., includes non-addressable devices from other processes.

devices_indices_map(global_shape)

Returns a mapping from devices to the array slices each contains.

The mapping includes all global devices, i.e., including non-addressable devices from other processes.

Parameters:

global_shape (Shape)

Return type:

Mapping[Device, Index]

is_equivalent_to(other, ndim)

Returns True if two shardings are equivalent.

Two shardings are equivalent if they place the same logical array shards on the same devices.

Parameters:

• self (Sharding)

• other (Sharding)

• ndim (int)

Return type:

bool

property is_fully_addressable: bool

Is this sharding fully addressable?

A sharding is fully addressable if the current process can address all of the devices named in the Sharding. is_fully_addressable is equivalent to “is_local” in multi-process JAX.

property is_fully_replicated: bool

Is this sharding fully replicated?

A sharding is fully replicated if each device has a complete copy of the entire data.

property memory_kind: str | None

Returns the memory kind of the sharding.

property num_devices: int

Number of devices that the sharding contains.

shard_shape(global_shape)

Returns the shape of the data on each device.

The shard shape returned by this function is calculated from global_shape and the properties of the sharding.

Parameters:

global_shape (Shape)

Return type:

Shape

with_memory_kind(kind)

Returns a new Sharding instance with the specified memory kind.

Parameters:

kind (str)

Return type:

Sharding

class jax.sharding.SingleDeviceSharding(*args, **kwargs)

Bases: Sharding

A Sharding that places its data on a single device.

Parameters:

device – A single Device.

Examples

>>> single_device_sharding = jax.sharding.SingleDeviceSharding(
...     jax.devices()[0])

property device_set: set[Device]

The set of devices that this Sharding spans.

In multi-controller JAX, the set of devices is global, i.e., includes non-addressable devices from other processes.

devices_indices_map(global_shape)

Returns a mapping from devices to the array slices each contains.

The mapping includes all global devices, i.e., including non-addressable devices from other processes.

Parameters:

global_shape (Shape)

Return type:

Mapping[Device, Index]

property is_fully_addressable: bool

Is this sharding fully addressable?

A sharding is fully addressable if the current process can address all of the devices named in the Sharding. is_fully_addressable is equivalent to “is_local” in multi-process JAX.

property is_fully_replicated: bool

Is this sharding fully replicated?

A sharding is fully replicated if each device has a complete copy of the entire data.

property memory_kind: str | None

Returns the memory kind of the sharding.

property num_devices: int

Number of devices that the sharding contains.

with_memory_kind(kind)

Returns a new Sharding instance with the specified memory kind.

Parameters:

kind (str)

Return type:

SingleDeviceSharding

class jax.sharding.NamedSharding(*args, **kwargs)

Bases: Sharding

A NamedSharding expresses sharding using named axes.

A NamedSharding is a pair of a Mesh of devices and PartitionSpec which describes how to shard an array across that mesh.

A Mesh is a multidimensional NumPy array of JAX devices, where each axis of the mesh has a name, e.g. 'x' or 'y'.

A PartitionSpec is a tuple, whose elements can be a None, a mesh axis, or a tuple of mesh axes. Each element describes how an input dimension is partitioned across zero or more mesh dimensions. For example, PartitionSpec('x', 'y') says that the first dimension of data is sharded across x axis of the mesh, and the second dimension is sharded across y axis of the mesh.

The Distributed arrays and automatic parallelization (https://docs.jax.dev/en/latest/notebooks/Distributed_arrays_and_automatic_parallelization.html#namedsharding-gives-a-way-to-express-shardings-with-names) tutorial has more details and diagrams that explain how Mesh and PartitionSpec are used.

Parameters:

• mesh – A jax.sharding.Mesh object.

• spec – A jax.sharding.PartitionSpec object.

Examples

>>> from jax.sharding import Mesh
>>> from jax.sharding import PartitionSpec as P
>>> mesh = Mesh(np.array(jax.devices()).reshape(2, 4), ('x', 'y'))
>>> spec = P('x', 'y')
>>> named_sharding = jax.sharding.NamedSharding(mesh, spec)

property addressable_devices: set[Device]

The set of devices in the Sharding that are addressable by the current process.

property device_set: set[Device]

The set of devices that this Sharding spans.

In multi-controller JAX, the set of devices is global, i.e., includes non-addressable devices from other processes.

property is_fully_addressable: bool

Is this sharding fully addressable?

A sharding is fully addressable if the current process can address all of the devices named in the Sharding. is_fully_addressable is equivalent to “is_local” in multi-process JAX.

property is_fully_replicated: bool

Is this sharding fully replicated?

A sharding is fully replicated if each device has a complete copy of the entire data.

property memory_kind: str | None

Returns the memory kind of the sharding.

property mesh

(self) -> object

property num_devices: int

Number of devices that the sharding contains.

property spec

(self) -> object

with_memory_kind(kind)

Returns a new Sharding instance with the specified memory kind.

Parameters:

kind (str)

Return type:

NamedSharding

class jax.sharding.PmapSharding(*args, **kwargs)

Bases: Sharding

Describes a sharding used by jax.pmap().

classmethod default(shape, sharded_dim=0, devices=None)

Creates a PmapSharding which matches the default placement used by jax.pmap().

Parameters:

• shape (Shape) – The shape of the input array.

• sharded_dim (int | None) – Dimension the input array is sharded on. Defaults to 0.

• devices (Sequence[xc.Device] | None) – Optional sequence of devices to use. If omitted, the implicit device order used by pmap is used, which is the order of jax.local_devices().

Return type:

PmapSharding

property device_set: set[Device]

The set of devices that this Sharding spans.

In multi-controller JAX, the set of devices is global, i.e., includes non-addressable devices from other processes.

property devices

(self) -> ndarray

devices_indices_map(global_shape)

Returns a mapping from devices to the array slices each contains.

The mapping includes all global devices, i.e., including non-addressable devices from other processes.

Parameters:

global_shape (Shape)

Return type:

Mapping[Device, Index]

is_equivalent_to(other, ndim)

Returns True if two shardings are equivalent.

Two shardings are equivalent if they place the same logical array shards on the same devices.

Parameters:

• self (PmapSharding)

• other (PmapSharding)

• ndim (int)

Return type:

bool

property is_fully_addressable: bool

Is this sharding fully addressable?

A sharding is fully addressable if the current process can address all of the devices named in the Sharding. is_fully_addressable is equivalent to “is_local” in multi-process JAX.

property is_fully_replicated: bool

Is this sharding fully replicated?

A sharding is fully replicated if each device has a complete copy of the entire data.

property memory_kind: str | None

Returns the memory kind of the sharding.

property num_devices: int

Number of devices that the sharding contains.

shard_shape(global_shape)

Returns the shape of the data on each device.

The shard shape returned by this function is calculated from global_shape and the properties of the sharding.

Parameters:

global_shape (Shape)

Return type:

Shape

property sharding_spec

(self) -> jax::ShardingSpec

with_memory_kind(kind)

Returns a new Sharding instance with the specified memory kind.

Parameters:

kind (str)

class jax.sharding.PartitionSpec(*args, **kwargs)

Tuple describing how to partition an array across a mesh of devices.

Each element is either None, a string, or a tuple of strings. See the documentation of jax.sharding.NamedSharding for more details.

This class exists so JAX’s pytree utilities can distinguish a partition specifications from tuples that should be treated as pytrees.

property reduced

(self) -> frozenset

property unreduced

(self) -> frozenset

class jax.sharding.Mesh(devices, axis_names, axis_types=None)

Declare the hardware resources available in the scope of this manager.

See the Distributed arrays and automatic parallelization tutorial (https://docs.jax.dev/en/latest/notebooks/Distributed_arrays_and_automatic_parallelization.html) and Explicit sharding tutorial (https://docs.jax.dev/en/latest/notebooks/explicit-sharding.html)

Parameters:

• devices (np.ndarray) – A NumPy ndarray object containing JAX device objects (as obtained e.g. from jax.devices()).

• axis_names (tuple[MeshAxisName, ...]) – A sequence of resource axis names to be assigned to the dimensions of the devices argument. Its length should match the rank of devices.

• axis_types (tuple[AxisType, ...] | None) – and optional tuple of jax.sharding.AxisType entries corresponding to the axis_names. See Explicit Sharding for more information.

Examples

>>> from jax.sharding import Mesh
>>> from jax.sharding import PartitionSpec as P, NamedSharding
>>> import numpy as np
...
>>> # Declare a 2D mesh with axes `x` and `y`.
>>> devices = np.array(jax.devices()).reshape(4, 2)
>>> mesh = Mesh(devices, ('x', 'y'))
>>> inp = np.arange(16).reshape(8, 2)
>>> arr = jax.device_put(inp, NamedSharding(mesh, P('x', 'y')))
>>> out = jax.jit(lambda x: x * 2)(arr)
>>> assert out.sharding == NamedSharding(mesh, P('x', 'y'))