(
nrep, axis_size, npart, parts, arg_shape, map_axis: Optional[int]
)
| 411 | |
| 412 | |
| 413 | def _pmap_sharding_spec( |
| 414 | nrep, axis_size, npart, parts, arg_shape, map_axis: Optional[int] |
| 415 | ) -> pmap_lib.ShardingSpec: |
| 416 | replication_factor, ragged = divmod(nrep, axis_size) |
| 417 | assert not ragged |
| 418 | # get the sharding spec from inner sharded_jits as if we weren't in a pmap |
| 419 | pspec = partitioned_sharding_spec(npart, parts, arg_shape) |
| 420 | maybe_replicate = ( |
| 421 | () if replication_factor == 1 else (pmap_lib.Replicated(replication_factor),) |
| 422 | ) |
| 423 | if map_axis is not None: |
| 424 | sharded_in_axis = sum( |
| 425 | not isinstance(s, pmap_lib.NoSharding) for s in pspec.sharding[:map_axis] |
| 426 | ) |
| 427 | |
| 428 | def shift_sharded_axis(a): |
| 429 | if isinstance(a, pmap_lib.ShardedAxis) and a.axis >= sharded_in_axis: |
| 430 | return pmap_lib.ShardedAxis(a.axis + 1) |
| 431 | return a |
| 432 | |
| 433 | # replication_factor represents the product of inner pmaps, so it goes |
| 434 | # after the outer pmapped axis at index 0 |
| 435 | return pmap_lib.ShardingSpec( |
| 436 | sharding=tuple_insert( |
| 437 | pspec.sharding, map_axis, pmap_lib.Unstacked(axis_size) |
| 438 | ), |
| 439 | mesh_mapping=it.chain( |
| 440 | [pmap_lib.ShardedAxis(sharded_in_axis)], |
| 441 | maybe_replicate, |
| 442 | map(shift_sharded_axis, pspec.mesh_mapping), |
| 443 | ), |
| 444 | ) |
| 445 | else: |
| 446 | return pmap_lib.ShardingSpec( |
| 447 | sharding=pspec.sharding, |
| 448 | mesh_mapping=(pmap_lib.Replicated(axis_size),) |
| 449 | + maybe_replicate |
| 450 | + pspec.mesh_mapping, |
| 451 | ) |
| 452 | |
| 453 | |
| 454 | def _get_pmap_sharding(devices, specs): |
no test coverage detected