Skip apply_gufunc compensation when post-rechunk block fits array.chunk-size

[thodson-usgs]

• Tests added
• Passes pre-commit run --all-files

Follow-up to #11683. The loop-dim compensation block there unconditionally shrinks loop chunks after rechunking a core dim to -1, even when the original block was already small. Result: task graphs explode without any memory benefit.

import dask.array as da

src = da.zeros((200, 400), chunks=(2, 400))           # 100 input chunks, ~6 KB each
out = da.apply_gufunc(
    lambda x, c: x[:1], "(i),(i)->(j)",
    src, da.arange(200.0),
    axes=[(0,), (0,), (0,)], output_sizes={"j": 50},
    allow_rechunk=True, output_dtypes=float,
)
len(out.__dask_graph__())
# before: 20,701       after: 404

Fix: guard the compensation with an array.chunk-size budget. Only shrink loop dims when the post-rechunk block would actually exceed the limit. The memory-protection branch from #11683 still fires when it should — new test test_gufunc_chunksizes_adjustment_above_limit forces the budget down and pins that behavior.

Related: pydata/xarray#9907 (the original report motivating #11683; already closed), pydata/xarray#10130 (open memory-OOM report where the compensation over-fire contributes to graph blowup).

Companion xarray PR (pydata/xarray#11312) addresses the architectural side for interp(method="linear"|"nearest"); this one helps every apply_gufunc(allow_rechunk=True) caller.

cc @phofl @crusaderky

[github-actions]

Unit Test Results

See test report for an extended history of previous test failures. This is useful for diagnosing flaky tests.

     21 files  ± 0       21 suites  ±0   5h 25m 50s ⏱️ + 9m 34s
 18 294 tests + 1   16 935 ✅ ± 0   1 272 💤 ±0  87 ❌ +1 
317 539 runs  +19  273 745 ✅ +16  43 707 💤 +2  87 ❌ +1 

For more details on these failures, see this check.

Results for commit e9e8dba. ± Comparison against base commit c6a85f3.

This pull request removes 1 and adds 2 tests. Note that renamed tests count towards both.

dask.array.tests.test_gufunc ‑ test_gufunc_chunksizes_adjustment

dask.array.tests.test_gufunc ‑ test_gufunc_chunksizes_adjustment_above_limit
dask.array.tests.test_gufunc ‑ test_gufunc_chunksizes_adjustment_below_limit

[thodson-usgs]

Closing this. After benchmarking, the regime this fix targets (post-rechunk blocks well below array.chunk-size) is narrow — at normal chunk sizes there's no measurable difference, and the specific case that motivated it (pydata/xarray#9907, #10130) is being addressed from the xarray side in pydata/xarray#11312. Happy to revisit if someone hits the same symptom through a different entry point.

[This is Claude Code on behalf of Tim Hodson]