code: def test_assign_shape_info():
      from torch.utils._pytree import tree_map
  
      cases = [
          (
              ShapeTensor(torch.randn(3)),
              MutableShape(3),
          ),
          (
              (
                  ShapeTensor(torch.randn(3, 4)),
                  ShapeTensor(torch.randn(4, 5)),
              ),
              (MutableShape(3, 4), MutableShape(4, 5)),
          ),
          (
              [()],
              [()],
          ),
          (
              ([],),
              ([],),
          ),
          (
              {"a": ()},
              {"a": ()},
          ),
          (
              {"a": 0, "b": [{"c": 1}]},
              {"a": 0, "b": [{"c": 1}]},
          ),
          (
              {
                  "a": 0,
                  "b": [1, {"c": 2}, ShapeTensor(torch.randn(3))],
                  "c": (ShapeTensor(torch.randn(2, 3)), 1),
              },
              {
                  "a": 0,
                  "b": [1, {"c": 2}, MutableShape(3)],
                  "c": (MutableShape(2, 3), 1),
              },
          ),
      ]
  
      for tensors, shapes in cases:
          _shape._assign_shape_info(tensors, shapes)
  
          def _check_shape(obj):
              if isinstance(obj, ShapeTensor):
                  assert obj.real_shape is not None
  
          tree_map(_check_shape, tensors)
  
imports: import torch
  import torch.nn.functional
  from torch import nn
  import pytest
  import nni
  import nni.nas.profiler.pytorch.utils.shape
  from nni.mutable import MutableList, Categorical
  from nni.nas.nn.pytorch import MutableLinear, MutableConv2d, MutableDropout, MutableMultiheadAttention, ModelSpace, LayerChoice, InputChoice, Repeat
  from nni.nas.profiler.pytorch.utils import shape_inference, submodule_input_output_shapes, ShapeTensor, MutableShape, profiler_leaf_module
label_desc: Use image node_id 2 for calling a global function with example usage: test_assign_shape_info() without return types