jax.make_array_from_process_local_data

jax.make_array_from_process_local_data(sharding, local_data, global_shape=None)

使用进程中可用的数据创建分布式张量。

此函数是 make_array_from_callback 的常见特例。它假定数据在进程中可用，并负责索引的转换。

最常见的情况是分片沿批次维度进行分片，每个主机仅加载其对应的子批次。此函数还支持更一般的情况，例如混合多主机和多轴复制和分片，但您需要正确计算进程本地数据的大小和内容以满足分片约束。

特别是，如果任何两个主机是副本，则 host_local_data 也应该相同。

global_shape 是可选的。如果未提供，将根据 local_data 和 sharding 推断，假设每个主机仅代表其自己的数据以实现均匀分片。如果分片不均匀（见下文说明），则会引发异常。

显式设置 global_shape 允许更精细地控制，并且可以与非均匀分片一起使用。global_shape 的每个维度必须匹配 host_local_data，或匹配 sharding 推断出的全局形状（在这种情况下，它等同于将其设置为 None，但更明确）。

例如，如果维度 i 被完全分片，则该大小将是 per_device_shape[i] * jax.local_device_count()。每个设备将被映射到 local_data 数组的本地切片。例如，如果给定的进程地址切片是 (8, 12) 和 (24, 28)，那么这些切片将被映射到 local_data 的 (0, 4) 和 (4, 8)。

对于 global_shapes 与 local_shape 匹配的每个维度，每个设备将查找 local_data 中的切片。例如，如果 global_shape == local_data.shape，则假定本地数据是将被分片到设备中的实际目标数组。

如果 global_shape 与 local_data.shape 相同，则所有主机的数据必须相同。

示例

>>> from jax.sharding import PartitionSpec as P
>>> mesh_rows = 2
>>> mesh_cols =  jax.device_count() // 2
...
>>> mesh = jax.sharding.Mesh(np.array(jax.devices()).reshape(mesh_rows, mesh_cols), ('x', 'y'))

>>> sharding = jax.sharding.NamedSharding(mesh, P(('x', 'y'),))
>>> rows_per_device = 2
>>> feature_length = 32
>>> per_device_shape = (rows_per_device, feature_length)
>>> per_host_shape = (rows_per_device * len(mesh.local_devices), feature_length)
>>> per_host_generator = lambda : np.arange(np.prod(per_host_shape)).reshape(per_host_shape)
>>> per_host_data = per_host_generator()  # replace with your own per-host data pipeline that outputs numpy arrays
>>> global_shape = (rows_per_device * len(sharding.device_set), ) + per_device_shape[1:]
>>> output_global_array = jax.make_array_from_process_local_data(sharding, per_host_data, global_shape)
...
>>> assert output_global_array.addressable_data(0).shape == per_device_shape
>>> assert output_global_array.shape == global_shape

注意：虽然大多数分片是均匀的，但有可能设计一个异域分片网格，其中每个进程的设备在某些维度上以非网格状模式排列，或者索引非平凡地重叠。这种分片在那些维度上称为“非均匀”分片。在这种情况下，那些方向上的全局形状必须与局部形状匹配，因为没有有意义的方式来表示所有需要的每个进程数据而不会重叠。例如，对于全局形状 4x4，如果分片看起来像这样

0123 2103 4675 4567

有 4 个进程，分别包含设备 (0,1)、(2, 3)、(4, 5)、(6, 7)。那么每个主机的数据看起来像

xx.. ..xx …. …. .xx. x..x …. …. …. …. x..x .xx. …. …. xx.. ..xx

分片在行上是均匀的（每个主机需要行 1-2 或行 3-4），在列上是非均匀的（主机需要重叠但不匹配的列集）。因此，即使每个主机可能适合 2x2 的形状，所有主机的 local 数据必须具有 2x4 或 4x4 的形状。在这种情况下，用户必须显式提供 global_shape，对于 local_shape=(2, 4)，可能的有效 global shapes 是 (2, 4) 和 (4, 4)。

另一方面，对于分片

0213 x.x. .x.x. …. …. 0213 x.x. .x.x. …. …. 4657 …. …. .x.x x.x. 4657 …. …. .x.x x.x.

对于 local_shape=(2, 2)，此函数可以接受 2x2、2x4、4x2 和 4x4 的全局形状。将 global_shape 设置为 None，在这种情况下等同于将其设置为 (4, 4)。

参数:

• sharding – 全局数组的分片。

• local_data – 主机上的数据，将放置在本地设备上。每个维度应匹配 global_shape，或匹配 num_addressable_indices(dim)。

• global_shape – 全局数组的目标形状。如果为 None，则从 local_data 和 sharding 推断。

返回:

将具有 sharding=sharding 且形状为 global_shape 的张量。