code: def tree_map_zip(fn: Any, *pytrees):
      assert len(pytrees) > 0
      if len(pytrees) == 1:
          return tree_map(fn, pytrees[0])
      else:
          flat_args_list, spec_list = [], []
          for pytree in pytrees:
              flat_args, spec = tree_flatten(pytree)
              flat_args_list.append(flat_args)
              spec_list.append(spec)
          assert all(
              len(args) == len(flat_args_list[0])
              for args in flat_args_list
          ), "Inconsistent tree nodes length."
          return tree_unflatten(
              [fn(*args) for args in zip(*flat_args_list)], spec_list[0]
          )
  
file_path: nni/nni/compression/speedup/utils.py
imports: from copy import deepcopy
  import logging
  from typing import Any, Dict, List, Tuple
  import torch
  from torch.fx.immutable_collections import immutable_dict, immutable_list
  from torch.utils._pytree import tree_map, tree_flatten, tree_unflatten, _register_pytree_node, Context
label_desc: Use image node_id 8 for calling a global function with example usage: tree_map_zip(fn) and returns: tree_map, tree_unflatten