Overlap with new_axis option is not trimmed correctly

[chourroutm]

Describe the issue:
I would like to expand a 3D array of chunk size (32,128,128) to a 4D array of chunk size (3,32,128,128) with a computation that relies on an overlap (chunk sizes (52,148,148) and (3,52,148,148)), then trimming the overlap (chunk size (3,32,128,128)).

I would prefer not to have to deal with the overlap myself in the gufunc.

Minimal Complete Verifiable Example:

import numpy as np
import dask, dask.array
import cupy as cp

def create_color(u):
    return np.stack([u*1.0,u*0.5,u*0.0])

x = np.ones((32, 128, 128))
y = dask.array.from_array(x)

overlap_depth = 10

colors_preview_overlap = dask.array.overlap.overlap(y.astype(np.float64),(overlap_depth, overlap_depth, overlap_depth),'reflect')
colors_preview_map = dask.array.core.map_blocks(create_color,colors_preview_overlap,meta=cp.array(()),chunks=(3,32+2*overlap_depth,128+2*overlap_depth,128+2*overlap_depth),new_axis=[0])
colors_preview_trimmed = dask.array.overlap.trim_overlap(colors_preview_map, (0,overlap_depth, overlap_depth, overlap_depth)) # same return with `trim_internal`

colors_preview_1 = colors_preview_trimmed.compute()
print(colors_preview_1.shape)

colors_preview_2 = y.astype(np.float64).map_overlap(create_color,overlap_depth,'reflect',meta=cp.array(()),chunks=(3,32+2*overlap_depth,128+2*overlap_depth,128+2*overlap_depth),new_axis=[0])
colors_preview_2 = colors_preview_2.compute()
print(colors_preview_2.shape)

colors_preview_3 = y.astype(np.float64).map_overlap(create_color,(overlap_depth, overlap_depth, overlap_depth),'reflect',meta=cp.array(()),chunks=(3,32+2*overlap_depth,128+2*overlap_depth,128+2*overlap_depth),new_axis=[0])
colors_preview_3 = colors_preview_3.compute()
print(colors_preview_3.shape)

colors_preview_4 = y.astype(np.float64).map_overlap(create_color,(0,overlap_depth, overlap_depth, overlap_depth),'reflect',meta=cp.array(()),chunks=(3,32+2*overlap_depth,128+2*overlap_depth,128+2*overlap_depth),new_axis=[0])
colors_preview_4 = colors_preview_4.compute()
print(colors_preview_4.shape)

These will return:

(3, 52, 148, 148)
(0, 32, 128, 148)
(0, 32, 128, 148)
(3, 12, 128, 148)

whereas I would like them to return:

(3, 32, 128, 128)

Anything else we need to know?:

No

Environment:

• Dask version: 2024.5.0
• Python version: Python 3.11.5
• Operating System: Windows 11
• Install method (conda, pip, source): pip

[dstansby]

Here's an even simpler reproducible example:

import dask.array as da
import numpy as np

x = da.arange(4, chunks=2)
print(x.compute())
print(x.map_blocks(lambda x: np.stack([x, x * 0.5]), new_axis=[0]).compute())
print(x.map_overlap(lambda x: np.stack([x, x * 0.5]), depth=1, new_axis=[0]).compute())

array([0, 1, 2, 3])
array([[0. , 1. , 2. , 3. ],
       [0. , 0.5, 1. , 1.5]])
array([[0. , 1. , 2. , 1. , 2. , 3. ],
       [0. , 0.5, 1. , 0.5, 1. , 1.5]])

It looks like the edge overlaps are being trimmed with map_overlap, but the interior overlaps are not being trimmed before returning the result.