jax.sharding
模块#
类#
- class jax.sharding.Sharding#
描述
jax.Array
如何跨设备布局。- addressable_devices_indices_map(global_shape)[源代码]#
从可寻址设备到每个设备包含的数组数据切片的映射。
addressable_devices_indices_map
包含应用于可寻址设备的device_indices_map
的部分。- 参数:
global_shape (Shape)
- 返回类型:
Mapping[Device, Index | None]
- devices_indices_map(global_shape)[源代码]#
返回从设备到每个设备包含的数组切片的映射。
映射包括所有全局设备,即包括来自其他进程的不可寻址设备。
- 参数:
global_shape (Shape)
- 返回类型:
Mapping[Device, Index]
- is_equivalent_to(other, ndim)[源代码]#
如果两个分片等效,则返回
True
。如果两个分片将相同的逻辑数组分片放置在相同的设备上,则它们是等效的。
例如,如果
NamedSharding
和PositionalSharding
都将数组的相同分片放置在相同的设备上,则它们可能是等效的。
- property is_fully_addressable: bool[源代码]#
此分片是否完全可寻址?
如果当前进程可以寻址
Sharding
中命名的所有设备,则分片是完全可寻址的。is_fully_addressable
在多进程 JAX 中等同于 “is_local”。
- class jax.sharding.SingleDeviceSharding#
基类:
Sharding
将其数据放置在单个设备上的
Sharding
。- 参数:
device – 单个
Device
。
示例
>>> single_device_sharding = jax.sharding.SingleDeviceSharding( ... jax.devices()[0])
- devices_indices_map(global_shape)[源代码]#
返回从设备到每个设备包含的数组切片的映射。
映射包括所有全局设备,即包括来自其他进程的不可寻址设备。
- 参数:
global_shape (Shape)
- 返回类型:
Mapping[Device, Index]
- property is_fully_addressable: bool[源代码]#
此分片是否完全可寻址?
如果当前进程可以寻址
Sharding
中命名的所有设备,则分片是完全可寻址的。is_fully_addressable
在多进程 JAX 中等同于 “is_local”。
- class jax.sharding.NamedSharding#
基类:
Sharding
NamedSharding
使用命名轴来表达分片。NamedSharding
是设备Mesh
和PartitionSpec
的一对,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) 教程包含更多详细信息和图表,解释了如何使用
Mesh
和PartitionSpec
。- 参数:
mesh –
jax.sharding.Mesh
对象。spec –
jax.sharding.PartitionSpec
对象。
示例
>>> 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 device_set: set[Device][source]#
此
Sharding
跨越的设备集合。在多控制器 JAX 中,设备集合是全局的,即包括来自其他进程的不可寻址设备。
- property is_fully_addressable: bool[source]#
此分片是否完全可寻址?
如果当前进程可以寻址
Sharding
中命名的所有设备,则分片是完全可寻址的。is_fully_addressable
在多进程 JAX 中等同于 “is_local”。
- property mesh#
(self) -> object
- property spec#
(self) -> object
- class jax.sharding.PositionalSharding(devices, *, memory_kind=None)[source]#
基类:
Sharding
- 参数:
devices (Sequence[xc.Device] | np.ndarray)
memory_kind (str | None)
- property is_fully_addressable: bool#
此分片是否完全可寻址?
如果当前进程可以寻址
Sharding
中命名的所有设备,则分片是完全可寻址的。is_fully_addressable
在多进程 JAX 中等同于 “is_local”。
- 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()
的顺序。
- 返回类型:
- property devices#
(self) -> ndarray
- devices_indices_map(global_shape)[source]#
返回从设备到每个设备包含的数组切片的映射。
映射包括所有全局设备,即包括来自其他进程的不可寻址设备。
- 参数:
global_shape (Shape)
- 返回类型:
Mapping[Device, Index]
- is_equivalent_to(other, ndim)[source]#
如果两个分片等效,则返回
True
。如果两个分片将相同的逻辑数组分片放置在相同的设备上,则它们是等效的。
例如,如果
NamedSharding
和PositionalSharding
都将数组的相同分片放置在相同的设备上,则它们可能是等效的。- 参数:
self (PmapSharding)
other (PmapSharding)
ndim (int)
- 返回类型:
- property is_fully_addressable: bool#
此分片是否完全可寻址?
如果当前进程可以寻址
Sharding
中命名的所有设备,则分片是完全可寻址的。is_fully_addressable
在多进程 JAX 中等同于 “is_local”。
- shard_shape(global_shape)[source]#
返回每个设备上数据的形状。
此函数返回的分片形状根据
global_shape
和分片的属性计算得出。- 参数:
global_shape (Shape)
- 返回类型:
Shape
- property sharding_spec#
(self) -> jax::ShardingSpec
- class jax.sharding.GSPMDSharding#
基类:
Sharding
- property is_fully_addressable: bool#
此分片是否完全可寻址?
如果当前进程可以寻址
Sharding
中命名的所有设备,则分片是完全可寻址的。is_fully_addressable
在多进程 JAX 中等同于 “is_local”。
- class jax.sharding.PartitionSpec(*partitions)[source]#
描述了如何在设备网格上分区数组的元组。
每个元素可以是
None
、字符串或字符串元组。 有关更多详细信息,请参阅jax.sharding.NamedSharding
的文档。这个类的存在是为了让 JAX 的 pytree 实用程序能够区分分区规范和应被视为 pytree 的元组。
- 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'))