Dealing with reduced dimensions in Dask `map_blocks`

dask_map_blocks_dims.py

# A method that adds a diagonal line:
def brighten(chunk):
    chunk_out = chunk[0] + np.eye(chunk[0].shape[0])
    return chunk_out[None,:5,:5]  # The key here is to use None or newaxis
  
# Make the array
arr = da.from_array(np.zeros((2,100,100)), chunks=(2,10,10))

# Examine the array shape
arr.shape
# >> (2, 100, 100)

# Use map_blocks to distribute:
op = da.map_blocks(brighten, arr, chunks=(1, 5,5)).compute()

# op looks like:
# array([[[1., 0., 0., ..., 0., 0., 0.],
#        [0., 1., 0., ..., 0., 0., 0.],
#        [0., 0., 1., ..., 1., 0., 0.],
#        ...,
#        [0., 0., 1., ..., 1., 0., 0.],
#        [0., 0., 0., ..., 0., 1., 0.],
#        [0., 0., 0., ..., 0., 0., 1.]]])