jax.sharding module

`jax.sharding` module#

Classes#

class jax.sharding.Sharding#

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

property addressable_devices: set[jaxlib.xla_extension.Device]#: The set of devices in the Sharding that are addressable by the current process.

addressable_devices_indices_map(global_shape)[source]#

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[jaxlib.xla_extension.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)[source]#

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 | None]

is_equivalent_to(other, ndim)[source]#

Returns True if two shardings are equivalent.

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

For example, a NamedSharding may be equivalent to a PositionalSharding if both place the same shards of the array on the same devices.

Parameters:

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.

shard_shape(global_shape)[source]#

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 (tuple[int, ...]) –
Return type:: tuple[int, ...]

with_memory_kind(kind)[source]#

Returns a new Sharding instance with the specified memory kind.

Parameters:: kind (str) –
Return type:: Sharding

class jax.sharding.XLACompatibleSharding#

Bases: Sharding

A Sharding that describes shardings expressible to XLA.

Subclasses of XLACompatibleSharding work with all JAX APIs and transformations that use XLA.

devices_indices_map(global_shape)[source]#

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 (tuple[int, ...]) –
Return type:: Mapping[Device, tuple[slice, ...]]

is_equivalent_to(other, ndim)[source]#

Returns True if two shardings are equivalent.

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

For example, a NamedSharding may be equivalent to a PositionalSharding if both place the same shards of the array on the same devices.

Parameters:

self (XLACompatibleSharding) –
other (XLACompatibleSharding) –
ndim (int) –

Return type:

bool

shard_shape(global_shape)[source]#

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 (tuple[int, ...]) –
Return type:: tuple[int, ...]

class jax.sharding.SingleDeviceSharding#

Bases: XLACompatibleSharding

A Sharding that places its data on a single device.

Parameters:: device – A single Device.

Example

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

property device_set: set[jaxlib.xla_extension.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)[source]#

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 (tuple[int, ...]) –
Return type:: Mapping[Device, tuple[slice, ...]]

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.

with_memory_kind(kind)[source]#

Returns a new Sharding instance with the specified memory kind.

Parameters:: kind (str) –
Return type:: SingleDeviceSharding

class jax.sharding.NamedSharding#

Bases: XLACompatibleSharding

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://jax.readthedocs.io/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.

Example

>>> 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[jaxlib.xla_extension.Device]#: The set of devices in the Sharding that are addressable by the current process.

property device_set: set[jaxlib.xla_extension.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.

with_memory_kind(kind)[source]#

Returns a new Sharding instance with the specified memory kind.

Parameters:: kind (str) –
Return type:: NamedSharding

class jax.sharding.PositionalSharding(devices, *, memory_kind=None)[source]#

Bases: XLACompatibleSharding

A sharding strategy that arranges data based on device positions.

This strategy enables efficient data distribution across devices, making it suitable for parallel processing on different hardware platforms.

Example

>>> devices = [xc.Device("GPU", 0), xc.Device("GPU", 1)] 
>>> sharding = PositionalSharding(devices) 
>>> print(sharding.shape)  # Output: (2,) 
>>> print(sharding.ndim)   # Output: 1 

Initialize a PositionalSharding instance with two GPU devices:

>>> devices = [xc.Device("GPU", 0), xc.Device("GPU", 1)] 
>>> sharding = PositionalSharding(devices) 

Parameters:

devices (Sequence[xc.Device] | np.ndarray) –
memory_kind (str | None) –

property T: PositionalSharding#

Create a new PositionalSharding instance with a transposed data layout.

Transposing a sharding instance involves rearranging the dimensions of the data layout. This method allows you to change the order of dimensions for efficient data processing or compatibility with other operations.

Parameters:: *axes – Axes permutation for the transpose operation as individual arguments.

Example

Transpose a PositionalSharding instance, swapping the first and second dimensions:

>>> transposed_sharding = sharding.transpose(1, 0) 

Returns:: A new PositionalSharding instance with the transposed data layout.
Return type:: PositionalSharding

property device_set: set[jaxlib.xla_extension.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.

replicate(axis=None, keepdims=True)[source]#

Create a new PositionalSharding instance with replicated data along a specified axis.

Replicating a sharding instance involves creating additional copies of the data to distribute it along a specific axis. This method allows you to replicate the data within the sharding strategy, which can be useful for parallel processing or data redundancy.

Parameters:

axis – The axis along which data replication is performed. If not specified, replication is applied across all dimensions.
keepdims – Whether to keep the dimensions of the sharding consistent or not. When set to True, the replicated dimensions are retained; otherwise, they are collapsed.

Example

Replicate a PositionalSharding instance along the first axis while keeping dimensions:

>>> replicated_sharding = sharding.replicate(axis=0, keepdims=True) 

Returns:: A new PositionalSharding instance with the data replicated along the specified axis and dimensionality adjustments based on the ‘keepdims’ parameter.
Return type:: PositionalSharding

reshape(*shape)[source]#

Returns a new PositionalSharding instance with a reshaped data layout.

Parameters:: *shape – New shape dimensions.
Returns:: A new PositionalSharding instance.
Return type:: PositionalSharding

Example

Reshape a PositionalSharding instance:

>>> new_sharding = sharding.reshape(2, 1) 

transpose(*axes)[source]#

Create a new PositionalSharding instance with a transposed data layout.

Parameters:: *axes – Axes permutation for the transpose operation as individual arguments.

Example

Transpose a PositionalSharding instance, swapping the first and second dimensions:

>>> transposed_sharding = sharding.transpose(1, 0) 

Returns:: A new PositionalSharding instance with the transposed data layout.
Return type:: PositionalSharding

with_memory_kind(kind)[source]#

Create a new PositionalSharding instance with a specified memory kind.

Memory kind refers to the type of memory or storage associated with the sharding strategy. This method allows you to customize the memory kind used by the sharding for efficient memory management and allocation.

Parameters:: kind (str) – The memory kind to associate with the sharding. Common memory kinds include “HBM” (High-Bandwidth Memory) and “DDR” (Double Data Rate memory).

Example

Create a new PositionalSharding instance with a specific memory kind, such as “HBM” (High-Bandwidth Memory): >>> sharding_with_mem_kind = sharding.with_memory_kind(“HBM”) # doctest: +SKIP

Returns:: A new PositionalSharding instance with the specified memory kind.
Return type:: PositionalSharding

class jax.sharding.PmapSharding#

Bases: XLACompatibleSharding

Describes a sharding used by jax.pmap().

classmethod default(shape, sharded_dim=0, devices=None)[source]#

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

Parameters:

shape (Shape) – The shape of the input array.
sharded_dim (int) – 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
used (device order used by pmap is) – jax.local_devices().
of (which is the order) – jax.local_devices().

Return type:

PmapSharding

property device_set: set[jaxlib.xla_extension.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)[source]#

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 (tuple[int, ...]) –
Return type:: Mapping[Device, tuple[slice, ...]]

is_equivalent_to(other, ndim)[source]#

Returns True if two shardings are equivalent.

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

For example, a NamedSharding may be equivalent to a PositionalSharding if both place the same shards of the array 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.

shard_shape(global_shape)[source]#

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 (tuple[int, ...]) –
Return type:: tuple[int, ...]

with_memory_kind(kind)[source]#

Returns a new Sharding instance with the specified memory kind.

Parameters:: kind (str) –

class jax.sharding.GSPMDSharding#

Bases: XLACompatibleSharding

property device_set: set[jaxlib.xla_extension.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)[source]#

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 (tuple[int, ...]) –
Return type:: Mapping[Device, tuple[slice, ...]]

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.

with_memory_kind(kind)[source]#

Returns a new Sharding instance with the specified memory kind.

Parameters:: kind (str) –
Return type:: GSPMDSharding

class jax.sharding.PartitionSpec(*partitions)[source]#

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.

class jax.sharding.Mesh(devices: np.ndarray | Sequence[xc.Device], axis_names: str | Sequence[MeshAxisName])[source]#

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

In particular, all axis_names become valid resource names inside the managed block and can be used e.g. in the in_axis_resources argument of jax.experimental.pjit.pjit(). Also see JAX’s multi-process programming model (https://jax.readthedocs.io/en/latest/multi_process.html) and the Distributed arrays and automatic parallelization tutorial (https://jax.readthedocs.io/en/latest/notebooks/Distributed_arrays_and_automatic_parallelization.html)

If you are compiling in multiple threads, make sure that the with Mesh context manager is inside the function that the threads will execute.

Parameters:

devices – A NumPy ndarray object containing JAX device objects (as obtained e.g. from jax.devices()).
axis_names – A sequence of resource axis names to be assigned to the dimensions of the devices argument. Its length should match the rank of devices.

Example

>>> from jax.experimental.pjit import pjit
>>> from jax.sharding import Mesh
>>> from jax.sharding import PartitionSpec as P
>>> import numpy as np
...
>>> inp = np.arange(16).reshape((8, 2))
>>> devices = np.array(jax.devices()).reshape(4, 2)
...
>>> # Declare a 2D mesh with axes `x` and `y`.
>>> global_mesh = Mesh(devices, ('x', 'y'))
>>> # Use the mesh object directly as a context manager.
>>> with global_mesh:
...   out = pjit(lambda x: x, in_shardings=None, out_shardings=None)(inp)

>>> # Initialize the Mesh and use the mesh as the context manager.
>>> with Mesh(devices, ('x', 'y')) as global_mesh:
...   out = pjit(lambda x: x, in_shardings=None, out_shardings=None)(inp)

>>> # Also you can use it as `with ... as ...`.
>>> global_mesh = Mesh(devices, ('x', 'y'))
>>> with global_mesh as m:
...   out = pjit(lambda x: x, in_shardings=None, out_shardings=None)(inp)

>>> # You can also use it as `with Mesh(...)`.
>>> with Mesh(devices, ('x', 'y')):
...   out = pjit(lambda x: x, in_shardings=None, out_shardings=None)(inp)

jax.sharding module

Contents

jax.sharding module#

Classes#

`jax.sharding` module#