code: def _assign_shape_info(tensors, shapes):
      """Assign shape information returned by formula, to tensors returned by module forward."""
      from torch.utils._pytree import tree_flatten, tree_unflatten
  
      tensors_flatten, tree_spec_tensor = tree_flatten(tensors)
      shapes_flatten, tree_spec_shapes = tree_flatten(shapes)
      if tree_spec_tensor != tree_spec_shapes:
          raise ValueError(
              "Shape formula must return the same structure as the output, "
              f"but got {tree_spec_tensor} vs {tree_spec_shapes}"
          )
      if len(tensors_flatten) != len(shapes_flatten):
          raise ValueError(
              "Shape formula must return the same number of elements as the number of outputs, "
              f"but got {tensors_flatten} vs {shapes_flatten}"
          )
      for i, s in enumerate(shapes_flatten):
          if isinstance(tensors_flatten[i], ShapeTensor):
              if isinstance(s, MutableShape):
                  tensors_flatten[i].real_shape = s
              elif isinstance(s, ShapeTensor):
                  tensors_flatten[i] = s
              else:
                  raise ValueError(
                      f"Expected MutableShape or ShapeTensor in return value of shape formula, but got {type(s)}"
                  )
  
      return tree_unflatten(tensors_flatten, tree_spec_tensor)
  
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 3 for calling a global function with example usage: _assign_shape_info(tensors, shapes) and returns: tree_unflatten