code: def fast_map_structure_with_path(func, *structure):
      """Faster map_structure_with_path implementation."""
      head_entries_with_path = tree.flatten_with_path(structure[0])
      if len(structure) > 1:
          tail_entries = (tree.flatten(s) for s in structure[1:])
          entries_with_path = [
              e[0] + e[1:]
              for e in zip(head_entries_with_path, *tail_entries)
          ]
      else:
          entries_with_path = head_entries_with_path
      # Arbitrarily choose one of the structures of the original sequence (the last)
      # to match the structure for the flattened sequence.
      return tree.unflatten_as(
          structure[-1], [func(*x) for x in entries_with_path]
      )
  
file_path: acme/acme/utils/tree_utils.py
imports: from typing import Sequence, List, TypeVar, Any
  import numpy
  import tree
label_desc: Use image node_id 2 for calling a global function with example usage: fast_map_structure_with_path(func) and returns: tree