map_overlap does not trim properly with new_axis when input array has more than 2 dimensions #11961

joshua-gould · 2025-05-21T13:46:37Z

Minimal Complete Verifiable Example:

import dask.array as da
import numpy as np


def process_chunk(x):
    return np.ones(shape=(2,) + x.shape)


# fails with 3d array
a = da.random.random((4, 100, 100)).rechunk((4, 10, 10))
result_shape = da.map_overlap(
    process_chunk,
    a,
    new_axis=0,
    depth={0: 0, 1: 2, 2: 2},
    trim=True,
    boundary="none",
).compute().shape
assert result_shape[1:] == a.shape, f"{result_shape} != {a.shape}"

# works with 2d array
a = da.random.random((100, 100)).rechunk((10, 10))
result_shape = da.map_overlap(
    process_chunk,
    a,
    new_axis=0,
    depth={0: 2, 1: 2},
    trim=True,
    boundary="none",
).compute().shape
assert result_shape[1:] == a.shape, f"{result_shape} != {a.shape}"

Environment:

Dask version: 2025.5.1
Python version: 3.12
Operating System: Linux
Install method: pip

github-actions bot added the needs triage Needs a response from a contributor label May 21, 2025

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

map_overlap does not trim properly with new_axis when input array has more than 2 dimensions #11961

map_overlap does not trim properly with new_axis when input array has more than 2 dimensions #11961

map_overlap does not trim properly with new_axis when input array has more than 2 dimensions #11961

map_overlap does not trim properly with new_axis when input array has more than 2 dimensions #11961

Comments