code: def create_pmap_sharding_spec(
      shape: tuple[int, ...],
      sharded_dim: int = 0,
      sharded_dim_size: int | None = None,
  ):
      if sharded_dim is not None:
          sharded_shape = shape[:sharded_dim] + shape[sharded_dim + 1 :]
          if sharded_dim_size is None:
              sharded_dim_size = shape[sharded_dim]
      else:
          assert sharded_dim_size is not None
          sharded_shape = shape
  
      return pmap_sharding_spec(
          sharded_dim_size, sharded_dim_size, sharded_shape, sharded_dim
      )
  
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 6 for calling a global function with example usage: create_pmap_sharding_spec(shape, sharded_dim, sharded_dim_size) and returns: pmap_sharding_spec