jax.sharding 模块#

#

class jax.sharding.Sharding#

描述 jax.Array 如何跨设备布局。

property addressable_devices: set[Device]#

Sharding 中可由当前进程寻址的设备集合。

addressable_devices_indices_map(global_shape)[源代码]#

从可寻址设备到每个设备包含的数组数据切片的映射。

addressable_devices_indices_map 包含应用于可寻址设备的 device_indices_map 的部分。

参数:

global_shape (Shape)

返回类型:

Mapping[Device, Index | None]

property device_set: set[Device][源代码]#

Sharding 跨越的设备集合。

在多控制器 JAX 中,设备集合是全局的,即包括来自其他进程的不可寻址设备。

devices_indices_map(global_shape)[源代码]#

返回从设备到每个设备包含的数组切片的映射。

映射包括所有全局设备,即包括来自其他进程的不可寻址设备。

参数:

global_shape (Shape)

返回类型:

Mapping[Device, Index]

is_equivalent_to(other, ndim)[源代码]#

如果两个分片等效,则返回 True

如果两个分片将相同的逻辑数组分片放置在相同的设备上,则它们是等效的。

例如,如果 NamedShardingPositionalSharding 都将数组的相同分片放置在相同的设备上,则它们可能是等效的。

参数:
返回类型:

bool

property is_fully_addressable: bool[源代码]#

此分片是否完全可寻址?

如果当前进程可以寻址 Sharding 中命名的所有设备,则分片是完全可寻址的。is_fully_addressable 在多进程 JAX 中等同于 “is_local”。

property is_fully_replicated: bool[源代码]#

此分片是否完全复制?

如果每个设备都具有完整数据的完整副本,则分片是完全复制的。

property memory_kind: str | None[源代码]#

返回分片的内存类型。

property num_devices: int[源代码]#

分片包含的设备数量。

shard_shape(global_shape)[源代码]#

返回每个设备上数据的形状。

此函数返回的分片形状根据 global_shape 和分片的属性计算得出。

参数:

global_shape (Shape)

返回类型:

Shape

with_memory_kind(kind)[源代码]#

返回具有指定内存类型的新 Sharding 实例。

参数:

kind (str)

返回类型:

Sharding

class jax.sharding.SingleDeviceSharding#

基类: Sharding

将其数据放置在单个设备上的 Sharding

参数:

device – 单个 Device

示例

>>> single_device_sharding = jax.sharding.SingleDeviceSharding(
...     jax.devices()[0])
property device_set: set[Device][源代码]#

Sharding 跨越的设备集合。

在多控制器 JAX 中,设备集合是全局的,即包括来自其他进程的不可寻址设备。

devices_indices_map(global_shape)[源代码]#

返回从设备到每个设备包含的数组切片的映射。

映射包括所有全局设备,即包括来自其他进程的不可寻址设备。

参数:

global_shape (Shape)

返回类型:

Mapping[Device, Index]

property is_fully_addressable: bool[源代码]#

此分片是否完全可寻址?

如果当前进程可以寻址 Sharding 中命名的所有设备,则分片是完全可寻址的。is_fully_addressable 在多进程 JAX 中等同于 “is_local”。

property is_fully_replicated: bool[源代码]#

此分片是否完全复制?

如果每个设备都具有完整数据的完整副本,则分片是完全复制的。

property memory_kind: str | None[源代码]#

返回分片的内存类型。

property num_devices: int[源代码]#

分片包含的设备数量。

with_memory_kind(kind)[源代码]#

返回具有指定内存类型的新 Sharding 实例。

参数:

kind (str)

返回类型:

SingleDeviceSharding

class jax.sharding.NamedSharding#

基类: Sharding

NamedSharding 使用命名轴来表达分片。

NamedSharding 是设备 MeshPartitionSpec 的一对,PartitionSpec 描述了如何跨该网格对数组进行分片。

Mesh 是 JAX 设备的多维 NumPy 数组,其中网格的每个轴都有一个名称,例如 'x''y'

PartitionSpec 是一个元组,其元素可以是 None、网格轴或网格轴元组。每个元素描述了如何跨零个或多个网格维度对输入维度进行分区。例如,PartitionSpec('x', 'y') 表示数据的第一个维度跨网格的 x 轴进行分片,第二个维度跨网格的 y 轴进行分片。

分布式数组和自动并行化 (https://jax.net.cn/en/latest/notebooks/Distributed_arrays_and_automatic_parallelization.html#namedsharding-gives-a-way-to-express-shardings-with-names) 教程包含更多详细信息和图表,解释了如何使用 MeshPartitionSpec

参数:

示例

>>> 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][source]#

Sharding 中可由当前进程寻址的设备集合。

property device_set: set[Device][source]#

Sharding 跨越的设备集合。

在多控制器 JAX 中,设备集合是全局的,即包括来自其他进程的不可寻址设备。

property is_fully_addressable: bool[source]#

此分片是否完全可寻址?

如果当前进程可以寻址 Sharding 中命名的所有设备,则分片是完全可寻址的。is_fully_addressable 在多进程 JAX 中等同于 “is_local”。

property is_fully_replicated: bool#

此分片是否完全复制?

如果每个设备都具有完整数据的完整副本,则分片是完全复制的。

property memory_kind: str | None[source]#

返回分片的内存类型。

property mesh#

(self) -> object

property num_devices: int[source]#

分片包含的设备数量。

property spec#

(self) -> object

with_memory_kind(kind)[source]#

返回具有指定内存类型的新 Sharding 实例。

参数:

kind (str)

返回类型:

NamedSharding

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

基类: Sharding

参数:
  • devices (Sequence[xc.Device] | np.ndarray)

  • memory_kind (str | None)

property device_set: set[xc.Device]#

Sharding 跨越的设备集合。

在多控制器 JAX 中,设备集合是全局的,即包括来自其他进程的不可寻址设备。

property is_fully_addressable: bool#

此分片是否完全可寻址?

如果当前进程可以寻址 Sharding 中命名的所有设备,则分片是完全可寻址的。is_fully_addressable 在多进程 JAX 中等同于 “is_local”。

property is_fully_replicated: bool#

此分片是否完全复制?

如果每个设备都具有完整数据的完整副本,则分片是完全复制的。

property memory_kind: str | None[source]#

返回分片的内存类型。

property num_devices: int[source]#

分片包含的设备数量。

with_memory_kind(kind)[source]#

返回具有指定内存类型的新 Sharding 实例。

参数:

kind (str)

返回类型:

PositionalSharding

class jax.sharding.PmapSharding#

基类: Sharding

描述了 jax.pmap() 使用的分片。

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

创建一个 PmapSharding,它匹配 jax.pmap() 使用的默认放置。

参数:
  • shape (Shape) – 输入数组的形状。

  • sharded_dim (int | None) – 输入数组被分片的维度。默认为 0。

  • devices (Sequence[xc.Device] | None | None) – 可选的设备序列以供使用。如果省略,则使用 pmap 使用的隐式设备顺序,即 jax.local_devices() 的顺序。

返回类型:

PmapSharding

property device_set: set[Device]#

Sharding 跨越的设备集合。

在多控制器 JAX 中,设备集合是全局的,即包括来自其他进程的不可寻址设备。

property devices#

(self) -> ndarray

devices_indices_map(global_shape)[source]#

返回从设备到每个设备包含的数组切片的映射。

映射包括所有全局设备,即包括来自其他进程的不可寻址设备。

参数:

global_shape (Shape)

返回类型:

Mapping[Device, Index]

is_equivalent_to(other, ndim)[source]#

如果两个分片等效,则返回 True

如果两个分片将相同的逻辑数组分片放置在相同的设备上,则它们是等效的。

例如,如果 NamedShardingPositionalSharding 都将数组的相同分片放置在相同的设备上,则它们可能是等效的。

参数:
返回类型:

bool

property is_fully_addressable: bool#

此分片是否完全可寻址?

如果当前进程可以寻址 Sharding 中命名的所有设备,则分片是完全可寻址的。is_fully_addressable 在多进程 JAX 中等同于 “is_local”。

property is_fully_replicated: bool#

此分片是否完全复制?

如果每个设备都具有完整数据的完整副本,则分片是完全复制的。

property memory_kind: str | None[source]#

返回分片的内存类型。

property num_devices: int[source]#

分片包含的设备数量。

shard_shape(global_shape)[source]#

返回每个设备上数据的形状。

此函数返回的分片形状根据 global_shape 和分片的属性计算得出。

参数:

global_shape (Shape)

返回类型:

Shape

property sharding_spec#

(self) -> jax::ShardingSpec

with_memory_kind(kind)[source]#

返回具有指定内存类型的新 Sharding 实例。

参数:

kind (str)

class jax.sharding.GSPMDSharding#

基类: Sharding

property device_set: set[Device]#

Sharding 跨越的设备集合。

在多控制器 JAX 中,设备集合是全局的,即包括来自其他进程的不可寻址设备。

property is_fully_addressable: bool#

此分片是否完全可寻址?

如果当前进程可以寻址 Sharding 中命名的所有设备,则分片是完全可寻址的。is_fully_addressable 在多进程 JAX 中等同于 “is_local”。

property is_fully_replicated: bool#

此分片是否完全复制?

如果每个设备都具有完整数据的完整副本,则分片是完全复制的。

property memory_kind: str | None[source]#

返回分片的内存类型。

property num_devices: int[source]#

分片包含的设备数量。

with_memory_kind(kind)[source]#

返回具有指定内存类型的新 Sharding 实例。

参数:

kind (str)

返回类型:

GSPMDSharding

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

描述了如何在设备网格上分区数组的元组。

每个元素可以是 None、字符串或字符串元组。 有关更多详细信息,请参阅 jax.sharding.NamedSharding 的文档。

这个类的存在是为了让 JAX 的 pytree 实用程序能够区分分区规范和应被视为 pytree 的元组。

index(value)[source]#

返回值的第一个索引。

如果值不存在,则引发 ValueError 错误。

class jax.sharding.Mesh(devices, axis_names, *, axis_types=None)[source]#

声明此管理器范围内可用的硬件资源。

请参阅分布式数组和自动并行化教程 (https://jax.net.cn/en/latest/notebooks/Distributed_arrays_and_automatic_parallelization.html) 和显式分片教程 (https://jax.net.cn/en/latest/notebooks/explicit-sharding.html)

参数:
  • devices (np.ndarray) – 包含 JAX 设备对象的 NumPy ndarray 对象(例如从 jax.devices() 获得)。

  • axis_names (tuple[MeshAxisName, ...]) – 要分配给 devices 参数维度的资源轴名称序列。其长度应与 devices 的秩匹配。

  • axis_types (tuple[AxisType, ...] | None)

示例

>>> 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'))