map_blocks raises AssertionError given chunks along a multiindex

What happened?

To parallelize a computation over a sparse array, I want to create chunks after stacking the array and dropping fill values. Such an array has a multiindex, which breaks the map_blocks method (with an unhelpful error). The AssertionError raised suggests to me that the method does not account for such an array having (coordinate) variables that are not dimensions or a scalar.

What did you expect to happen?

I expect the MCVE below to return its input unmodified.

Minimal Complete Verifiable Example

import xarray as xr

a = xr.DataArray([[0, 1], [2, 3]], {"x": [0, 1], "y": [0, 1]})
a = a.stack({"n": ("x", "y")})
a = a.chunk({"n": 2})
a.map_blocks(lambda x: x)

MVCE confirmation

• Minimal example — the example is as focused as reasonably possible to demonstrate the underlying issue in xarray.
• Complete example — the example is self-contained, including all data and the text of any traceback.
• Verifiable example — the example copy & pastes into an IPython prompt or Binder notebook, returning the result.
• New issue — a search of GitHub Issues suggests this is not a duplicate.
• Recent environment — the issue occurs with the latest version of xarray and its dependencies.

Relevant log output

File ~/tmp/bug/venv/lib/python3.11/site-packages/xarray/core/dataarray.py:5526, in DataArray.map_blocks(self, func, args, kwargs, template)
   5428 """
   5429 Apply a function to each block of this DataArray.
   5430 
   (...)
   5522     month    (time) int64 dask.array<chunksize=(24,), meta=np.ndarray>
   5523 """
   5524 from xarray.core.parallel import map_blocks
-> 5526 return map_blocks(func, self, args, kwargs, template)

File ~/tmp/bug/venv/lib/python3.11/site-packages/xarray/core/parallel.py:539, in map_blocks(func, obj, args, kwargs, template)
    535 for chunk_tuple in itertools.product(*ichunk.values()):
    536     # mapping from dimension name to chunk index
    537     chunk_index = dict(zip(ichunk.keys(), chunk_tuple))
--> 539     blocked_args = [
    540         subset_dataset_to_block(graph, gname, arg, input_chunk_bounds, chunk_index)
    541         if isxr
    542         else arg
    543         for isxr, arg in zip(is_xarray, npargs)
    544     ]
    546     # raise nice error messages in _wrapper
    547     expected: ExpectedDict = {
    548         # input chunk 0 along a dimension maps to output chunk 0 along the same dimension
    549         # even if length of dimension is changed by the applied function
   (...)
    563         },
    564     }

File ~/tmp/bug/venv/lib/python3.11/site-packages/xarray/core/parallel.py:540, in <listcomp>(.0)
    535 for chunk_tuple in itertools.product(*ichunk.values()):
    536     # mapping from dimension name to chunk index
    537     chunk_index = dict(zip(ichunk.keys(), chunk_tuple))
    539     blocked_args = [
--> 540         subset_dataset_to_block(graph, gname, arg, input_chunk_bounds, chunk_index)
    541         if isxr
    542         else arg
    543         for isxr, arg in zip(is_xarray, npargs)
    544     ]
    546     # raise nice error messages in _wrapper
    547     expected: ExpectedDict = {
    548         # input chunk 0 along a dimension maps to output chunk 0 along the same dimension
    549         # even if length of dimension is changed by the applied function
   (...)
    563         },
    564     }

File ~/tmp/bug/venv/lib/python3.11/site-packages/xarray/core/parallel.py:195, in subset_dataset_to_block(graph, gname, dataset, input_chunk_bounds, chunk_index)
    190     graph[chunk_variable_task] = (
    191         tuple,
    192         [variable.dims, chunk, variable.attrs],
    193     )
    194 else:
--> 195     assert name in dataset.dims or variable.ndim == 0
    197     # non-dask array possibly with dimensions chunked on other variables
    198     # index into variable appropriately
    199     subsetter = {
    200         dim: _get_chunk_slicer(dim, chunk_index, input_chunk_bounds)
    201         for dim in variable.dims
    202     }

AssertionError: 

Anything else we need to know?

No response

Environment

Start in xarray/core/parallel.py at subset_dataset_to_block, where the AssertionError occurs for the stacked DataArray's coordinate variables. Reproduce the issue with the provided map_blocks example and trace how the multiindex variables are handled. Done means the example returns its input unmodified without the assertion.