def map_parallel(
    fn: Callable[..., _T],
    *args: Collection[Any],
    timeout: float | None = None,
    max_workers: int | None = None,
) -> Sequence[_T]:
  """Runs `map(*args)` in parallel.

  IMPORTANT: Passed callables must be threadsafe.

  Args:
    fn: function to execute (MUST BE THREADSAFE)
    *args: arguments to pass to function.
    timeout: the maximum number of seconds to wait.
    max_workers: them maximum number of parallel jobs. If None, will use as many
      workers as there are arguments.

  Returns:
    [fn(arg0, arg1, ...) for arg0, arg1, ...  in zip(*args)]
    However, the calls will be executed concurrently.

  Raises:
    TimeoutError: If all the results are not generated before the timeout.
    Exception: If fn(*args) raises for any values.
  """
  tasks = {
      str(n): functools.partial(fn, *arg)
      for n, arg in enumerate(zip(*args, strict=True))
  }
  results = run_tasks(tasks, timeout=timeout, max_workers=max_workers)
  return [results[key] for key in tasks]
Runs `map(*args)` in parallel.

  IMPORTANT: Passed callables must be threadsafe.

  Args:
    fn: function to execute (MUST BE THREADSAFE)
    *args: arguments to pass to function.
    timeout: the maximum number of seconds to wait.
    max_workers: them maximum number of parallel jobs. If None, will use as many
      workers as there are arguments.

  Returns:
    [fn(arg0, arg1, ...) for arg0, arg1, ...  in zip(*args)]
    However, the calls will be executed concurrently.

  Raises:
    TimeoutError: If all the results are not generated before the timeout.
    Exception: If fn(*args) raises for any values.
  
concordia/utils/concurrency.py