code: def switch_case_shape_info(
      indicator: Any, branches: dict[str, Any]
  ) -> Any:
      """Combine shape information obtained from multiple branches together.
  
      The return type will look at one of the branches,
      ``switch_case`` will deeply merged into every shape dimension which is different across branches.
  
      Examples
      --------
      >>> switch_case_shape_info(nni.choice('x', [1, 2, 3]), {
      ...     1: MutableShape(2, 3),
      ...     2: MutableShape(4, 5),
      ...     3: MutableShape(6, 7),
      ... })
      MutableShape(
          switch_case(Categorical([1, 2, 3], label='x'), {1: 2, 2: 4, 3: 6}),
          switch_case(Categorical([1, 2, 3], label='x'), {1: 3, 2: 5, 3: 7})
      )
      """
      from torch.utils._pytree import tree_flatten, tree_unflatten
  
      cases, shapes_flattened, spec = [], [], None
      for case, item in branches.items():
          s_flattened, cur_spec = tree_flatten(item)
          if spec is None:
              spec = cur_spec
          elif spec != cur_spec:
              raise ValueError(
                  "Shape information must have the same structure, "
                  f"but got {spec} vs {cur_spec}"
              )
          if shapes_flattened and len(shapes_flattened[0]) != len(
              s_flattened
          ):
              raise ValueError(
                  "Shape information must have the same number of elements, "
                  f"but got {shapes_flattened[0]} vs {s_flattened}"
              )
          cases.append(case)
          shapes_flattened.append(s_flattened)
  
      assert spec is not None, "branches can not be empty"
  
      shapes_combined = []
  
      # This loop iterate over shapes of multiple tensors.
      for i in range(len(shapes_flattened[0])):
          # Shapes at this position are from different branches.
          to_merge = [s[i] for s in shapes_flattened]
  
          if not isinstance(to_merge[0], MutableShape):
              # Do nothing for non-MutableShape.
              shapes_combined.append(to_merge[0])
          else:
              # We make sure every shape to merge has the same dimension.
              assert all(
                  len(shape) == len(to_merge[0]) for shape in to_merge
              )
              # This loop iterate over elements of a shape.
              # Create every element of the output shape separately, so that the return value is a MutableShape.
              elements = []
              for j in range(len(to_merge[0])):
                  to_merge_j = [s[j] for s in to_merge]
                  if all(
                      _mutable_equal(s, to_merge_j[0])
                      for s in to_merge_j
                  ):
                      # If all elements are the same, we can just use the first one.
                      # This is to avoid creating an unnecessary MutableExpression.
                      elements.append(to_merge_j[0])
                  else:
                      elements.append(
                          MutableExpression.switch_case(
                              indicator,
                              {
                                  case: s
                                  for case, s in zip(cases, to_merge_j)
                              },
                          )
                      )
              shapes_combined.append(MutableShape(*elements))
  
      return tree_unflatten(shapes_combined, spec)
  
imports: from __future__ import annotations
  import functools
  import logging
  import traceback
  from typing import Union, Tuple, Callable, Iterable, Any, NamedTuple, Dict, cast, overload
  import torch
  from torch import nn
  from torch.utils.hooks import RemovableHandle
  from nni.mutable import MutableExpression, Mutable, LabeledMutable, Sample, SampleValidationError
  from nni.mutable.mutable import _mutable_equal
  from .misc import is_leaf_module, argument_in_spec
label_desc: Use image node_id 2 for calling a global function with example usage: switch_case_shape_info(indicator, branches) and returns: tree_unflatten