code: def _sharding_spec_mesh_shape(self):
      sharded_axis_sizes = []
      for sharding in self.sharding:
          if isinstance(sharding, NoSharding):
              continue
          elif isinstance(sharding, Unstacked):
              sharded_axis_sizes.append(sharding.size)
          elif isinstance(sharding, Chunked):
              sharded_axis_sizes.extend(sharding.chunks)
          else:
              util.assert_unreachable(sharding)
      return tuple(
          sharded_axis_sizes[a.axis]
          if isinstance(a, ShardedAxis)
          else a.replicas
          for a in self.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 1 for calling a global function with example usage: _sharding_spec_mesh_shape() and returns: tuple