code: def parallel_map(
      map_fn: Callable[
          ..., _Tout
      ],  # Callable[[_Tin0, _Tin1,...], Tout]
      *trees: Tree[_Tin],  # _Tin0, _Tin1,...
      max_workers: int = 32,
      report_progress: bool = False,
  ) -> Tree[_Tout]:  # pytype: disable=invalid-annotation
      """Same as `map_structure` but apply map_fn in parallel.
  
      Args:
        map_fn: Worker function
        *trees: Nested input to pass to the `map_fn`
        max_workers: Number of workers
        report_progress: If True, display a progression bar.
  
      Returns:
        The nested structure after `map_fn` has been applied.
      """
      with concurrent.futures.ThreadPoolExecutor(
          max_workers=max_workers,
      ) as executor:
          launch_worker = functools.partial(executor.submit, map_fn)
          futures = map_structure(launch_worker, *trees)
  
          leaves = tree.flatten(futures)
  
          itr = concurrent.futures.as_completed(leaves)
          if report_progress:
              itr = tqdm.tqdm(itr, total=len(leaves))
  
          for f in itr:  # Propagate exception to main thread.
              if f.exception():
                  raise f.exception()
  
      return map_structure(lambda f: f.result(), futures)
  
imports: import concurrent.futures
  import functools
  from typing import Callable, TypeVar
  from tensorflow_datasets.core.utils import tqdm_utils
  from tensorflow_datasets.core.utils import type_utils
  import tree
label_desc: Use image node_id 1 for calling a global function with example usage: parallel_map(map_fn) and returns: map_structure