code: def pmap_sharding_spec(
      nrep,
      axis_size,
      sharded_shape: Sequence[int],
      map_axis: int | None,
  ) -> ShardingSpec:
      """Sharding spec for arguments or results of a pmap.
      Args:
        nrep: number of local XLA replicas (product of local axis sizes)
        axis_size: local axis size for outer pmap
        sharded_aval: the aval of the value inside the outer pmap, an instance of
          a ShapedArray.
        map_axis: the axis along which the value is mapped in the outer pmap
      Returns:
        A ShardingSpec.
      """
      replication_factor, ragged = divmod(nrep, axis_size)
      assert not ragged
      pspec = ShardingSpec(
          sharding=[_UNSHARDED_INSTANCE] * len(sharded_shape),
          mesh_mapping=(),
      )
      maybe_replicate = (
          ()
          if replication_factor == 1
          else (Replicated(replication_factor),)
      )
      if map_axis is not None:
          sharded_in_axis = sum(
              not isinstance(s, NoSharding)
              for s in pspec.sharding[:map_axis]
          )
  
          def shift_sharded_axis(a: MeshDimAssignment):
              if (
                  isinstance(a, ShardedAxis)
                  and a.axis >= sharded_in_axis
              ):
                  return ShardedAxis(a.axis + 1)
              return a
  
          # replication_factor represents the product of inner pmaps, so it goes
          # after the outer pmapped axis at index 0
          return ShardingSpec(
              sharding=util.tuple_insert(
                  pspec.sharding, map_axis, Unstacked(axis_size)
              ),
              mesh_mapping=itertools.chain(
                  [ShardedAxis(sharded_in_axis)],
                  maybe_replicate,
                  map(shift_sharded_axis, pspec.mesh_mapping),
              ),
          )
      else:
          return ShardingSpec(
              sharding=pspec.sharding,
              mesh_mapping=(Replicated(axis_size),)
              + maybe_replicate
              + pspec.mesh_mapping,
          )
  
file_path: jax/jax/_src/sharding_specs.py
imports: from __future__ import annotations
  from collections.abc import Sequence
  import itertools
  import math
  from typing import Union
  import numpy
  from jax._src import util
  from jax._src.lib import pmap_lib
label_desc: Use image node_id 5 for calling a global function with example usage: pmap_sharding_spec(nrep, axis_size, sharded_shape, map_axis) and returns: ShardingSpec, ShardingSpec, a, ShardedAxis