jax.experimental.shard_map.shard_map

jax.experimental.shard_map.shard_map(f, mesh, in_specs, out_specs, check_rep=True, auto=frozenset({}))

对数据的碎片应用一个函数。

备注

shard_map 是一个实验性 API，仍可能发生变化。有关分片数据的介绍，请参阅 分片计算。要更深入地了解如何使用 shard_map，请参阅 使用 shard_map 进行 SPMD 多设备并行。

参数:

• f (Callable) – 可调用的映射对象。每次应用 f，或 f 的“实例”，以映射参数的分片作为输入，并生成输出分片。

• mesh (Mesh | AbstractMesh) – 一个 jax.sharding.Mesh 表示要在其上分片数据并在其上执行 f 实例的设备阵列。Mesh 的名称可以在 f 中的集体通信操作中使用。这通常由一个实用函数创建，例如 jax.experimental.mesh_utils.create_device_mesh()。

• in_specs (Specs) – 一个带有 PartitionSpec 实例的 pytree，其树结构是将被映射的 args 元组的树前缀。类似于 NamedSharding，每个 PartitionSpec 表示相应的参数（或参数子树）应如何沿 mesh 的命名轴进行分片。在每个 PartitionSpec 中，在某个位置提到 mesh 轴名称表示沿该位置轴对相应的参数数组轴进行分片；未提到轴名称表示复制。如果某个参数或参数子树对应的 spec 为 None，则该参数不分片。

• out_specs (Specs) – 一个包含 PartitionSpec 实例的 pytree，其树结构是 f 输出的树前缀。每个 PartitionSpec 表示相应的输出分片应如何连接。在每个 PartitionSpec 中，提及 mesh 轴名称的位置表示沿相应位置轴连接该网格轴的分片。未提及 mesh 轴名称表示承诺输出值在该网格轴上是相等的，并且应生成单个值而不是连接。

• check_rep (bool) – 如果为 True（默认），启用额外的有效性检查和自动微分优化。有效性检查涉及未在 out_specs 中提及的任何网格轴名称是否与 f 的输出复制方式一致。如果在 f 中使用 Pallas 内核，则必须设置为 False。

• auto (frozenset[AxisName]) – (实验性) 一组可选的轴名称，来自 mesh ，我们不会在这些轴上分片数据或映射函数，而是允许编译器控制分片。这些名称不能在 in_specs 、 out_specs 或 f 中的通信集合中使用。

返回:

一个可调用对象，它根据 mesh 和 in_specs 将输入函数 f 应用于分片数据。

示例

有关示例，请参阅 分片计算 或 使用 shard_map 进行 SPMD 多设备并行。