code: def mars_scheduler(
      dsk: dict, keys: Union[List[List[str]], List[str]]
  ):
      """
      A Dask-Mars scheduler
  
      This scheduler is intended to be compatible with existing
      dask user interface, no callbacks are implemented.
  
      Parameters
      ----------
      dsk: Dict
          Dask graph, represented as a task DAG dictionary.
      keys: Union[List[List[str]], List[str]]
          1d or 2d list of Dask graph keys whose values we wish to compute and return.
  
      Returns
      -------
      Object
          Computed values corresponding to the provided keys with same dimension.
      """
  
      if isinstance(keys, List) and not isinstance(
          keys[0], List
      ):  # 1d keys
          task = execute(mars_dask_get(dsk, keys))
          if not isinstance(task, List):
              task = [task]
          return map(lambda x: x.fetch(), task)
      else:  # 2d keys
          res = execute(reduce(mars_dask_get(dsk, keys))).fetch()
          if not isinstance(res, List):
              return [[res]]
          else:
              return res
  
imports: from typing import List, Tuple, Union
  from dask.core import istask, ishashable
  from ...remote import spawn
  from ...session import execute
  from .utils import reduce
label_desc: Use image node_id 1 for calling a global function with example usage: mars_scheduler(dsk, keys) and returns: map, list, res