code: def dim_squeeze(shape: Shape, dim: Optional[int] = None) -> DimMap:
      # FIXME: this is wrong when dim=None and one of the dimensions
      # equals size of the mesh. For example squeeze(DTensor(tensor(4), Shard[0])) could
      # end up as squeeze(tensor(1)) if we have 4 devices; this would lead to
      # removal of a dimension that is not actually a singleton.
      return tuple(
          InputDim(i)
          for i, s in enumerate(shape)
          if s > 1
          or (dim is not None and i != normalize_dim(dim, len(shape)))
      )
  
imports: from dataclasses import dataclass
  from typing import Callable, cast, Dict, Iterable, Optional, Sequence, Set, Tuple, Union
  import torch
  from torch import Tensor
  from torch._subclasses.fake_tensor import unset_fake_temporarily
  from torch.distributed._tensor._utils import compute_local_shape
  from torch.distributed._tensor.api import Shard
  from torch.distributed._tensor.op_schema import OpSchema, OutputSharding, RuntimeSchemaInfo
  from torch.distributed._tensor.ops.utils import normalize_dim, normalize_dims, prod, register_prop_rule
  from torch.distributed._tensor.placement_types import DTensorSpec, Placement, Replicate
  from torch.fx.experimental.proxy_tensor import disable_proxy_modes_tracing
label_desc: Use image node_id 12 for calling a global function with example usage: dim_squeeze(shape, dim) and returns: tuple