jax.sharding
模块#
类#
- class jax.sharding.Sharding#
描述
jax.Array
如何在设备之间分布。- addressable_devices_indices_map(global_shape)[source]#
可寻址设备到每个设备包含的数组数据切片的映射。
addressable_devices_indices_map
包含device_indices_map
中适用于可寻址设备的部分。- 参数:
global_shape (Shape)
- 返回类型:
Mapping[Device, Index | None]
- devices_indices_map(global_shape)[source]#
返回设备到每个设备包含的数组切片的映射。
此映射包括所有全局设备,即包括来自其他进程的不可寻址设备。
- 参数:
global_shape (Shape)
- 返回类型:
Mapping[Device, Index]
- is_equivalent_to(other, ndim)[source]#
如果两个分片等效,则返回
True
。如果两个分片将相同的逻辑数组分片放置在相同的设备上,则它们等效。
例如,
NamedSharding
可能等效于PositionalSharding
,如果两者都将数组的相同分片放置在相同的设备上。
- property is_fully_addressable: bool[source]#
此分片是否完全可寻址?
如果当前进程可以寻址
Sharding
中命名的所有设备,则分片是完全可寻址的。is_fully_addressable
等效于多进程 JAX 中的“is_local”。
- class jax.sharding.SingleDeviceSharding#
Bases:
Sharding
将数据放置在单个设备上的
Sharding
。- 参数:
device – 单个
Device
。
示例
>>> single_device_sharding = jax.sharding.SingleDeviceSharding( ... jax.devices()[0])
- devices_indices_map(global_shape)[source]#
返回设备到每个设备包含的数组切片的映射。
此映射包括所有全局设备,即包括来自其他进程的不可寻址设备。
- 参数:
global_shape (Shape)
- 返回类型:
Mapping[Device, Index]
- property is_fully_addressable: bool[source]#
此分片是否完全可寻址?
如果当前进程可以寻址
Sharding
中命名的所有设备,则分片是完全可寻址的。is_fully_addressable
等效于多进程 JAX 中的“is_local”。
- class jax.sharding.NamedSharding#
Bases:
Sharding
一个
NamedSharding
使用命名轴来表示分片。一个
NamedSharding
是一个Mesh
设备和PartitionSpec
的对,描述了如何在该网格上对数组进行分片。一个
Mesh
是 JAX 设备的多维 NumPy 数组,网格的每个轴都有一个名称,例如'x'
或'y'
。一个
PartitionSpec
是一个元组,其元素可以是None
,网格轴或网格轴的元组。每个元素描述了输入维度如何在零个或多个网格维度上进行分区。例如,PartitionSpec('x', 'y')
表示数据的第一个维度在网格的x
轴上进行分片,第二个维度在网格的y
轴上进行分片。分布式数组和自动并行化 (https://jax.ac.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 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]#
Bases:
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#
Bases:
Sharding
描述
jax.pmap()
使用的分片。- classmethod default(shape, sharded_dim=0, devices=None)[source]#
创建一个
PmapSharding
,它与jax.pmap()
使用的默认放置方式匹配。- 参数:
shape (Shape) – 输入数组的形状。
sharded_dim (int) – 输入数组被分片的维度。默认为 0。
devices (Sequence[xc.Device] | None | None) – 可选的设备序列。如果省略,则隐式
used (pmap 使用的设备顺序为) –
jax.local_devices()
.of (这是顺序) –
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#
Bases:
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)[source]#
声明此管理器范围内的可用硬件资源。
特别地,所有
axis_names
在受管块内部变为有效的资源名称,并且可以使用,例如在jax.experimental.pjit.pjit()
的in_axis_resources
参数中。另请参阅 JAX 的多进程编程模型 (https://jax.ac.cn/en/latest/multi_process.html) 和分布式数组和自动并行化教程 (https://jax.ac.cn/en/latest/notebooks/Distributed_arrays_and_automatic_parallelization.html)如果在多个线程中进行编译,请确保
with Mesh
上下文管理器位于线程将执行的函数内部。- 参数:
devices (np.ndarray) – 包含 JAX 设备对象(例如从
jax.devices()
获取)的 NumPy ndarray 对象。axis_names (tuple[MeshAxisName, ...]) – 要分配给
devices
参数维度的资源轴名称序列。其长度应与devices
的秩匹配。
示例
>>> 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)