code: def _fetch_dims(tree):
      shapes = []
      tree_type = type(tree)
      if tree_type is dict:
          for v in tree.values():
              shapes.extend(_fetch_dims(v))
      elif tree_type is list or tree_type is tuple:
          for t in tree:
              shapes.extend(_fetch_dims(t))
      elif tree_type is torch.Tensor:
          shapes.append(tree.shape)
      else:
          raise ValueError("Not supported")
  
      return shapes
  
file_path: H2O/h2o_flexgen/benchmark/third_party/transformers/src/transformers/models/esm/openfold_utils/chunk_utils.py
imports: import logging
  import math
  from functools import partial
  from typing import Any, Callable, Dict, Optional, Sequence, Tuple
  import torch
  from .tensor_utils import tensor_tree_map, tree_map
label_desc: Use image node_id 1 for calling a global function with example usage: _fetch_dims(tree) and returns: shapes