code: def tree_map(fn, tree, leaf_type):
      if isinstance(tree, dict):
          return dict_map(fn, tree, leaf_type)
      elif isinstance(tree, list):
          return [tree_map(fn, x, leaf_type) for x in tree]
      elif isinstance(tree, tuple):
          return tuple([tree_map(fn, x, leaf_type) for x in tree])
      elif isinstance(tree, leaf_type):
          return fn(tree)
      else:
          print(type(tree))
          raise ValueError("Not supported")
  
file_path: H2O/h2o_flexgen/benchmark/third_party/transformers/src/transformers/models/esm/openfold_utils/tensor_utils.py
imports: from functools import partial
  from typing import List
  import torch
  import torch.nn
label_desc: Use image node_id 10 for calling a global function with example usage: tree_map(fn, tree, leaf_type) and returns: dict_map, unknown, tuple, fn