Skip to content

Align convolution dask output dtype with numpy/cupy (#3618)#3622

Merged
brendancol merged 4 commits into
xarray-contrib:mainfrom
brendancol:deep-sweep-metadata-convolution-2026-07-02
Jul 6, 2026
Merged

Align convolution dask output dtype with numpy/cupy (#3618)#3622
brendancol merged 4 commits into
xarray-contrib:mainfrom
brendancol:deep-sweep-metadata-convolution-2026-07-02

Conversation

@brendancol

Copy link
Copy Markdown
Contributor

Closes #3618.

The dask backends of convolution_2d / convolve_2d declared float64
via meta=np.array(()) even when the eager numpy and cupy backends promote
an int or float32 raster to float32. So result.dtype changed depending on
the backend, and the declared dask dtype did not even match the float32
chunks dask computes.

Change

  • _convolve_2d_dask_numpy and _convolve_2d_dask_cupy now pass the
    promoted float dtype into the map_overlap meta, so the declared dtype
    matches both the eager backends and the computed chunks.

Backends

numpy, cupy, dask+numpy, dask+cupy all verified on a CUDA host: for a
float32 input every backend now reports float32 (was float64 on the two
dask paths). float64 input is unchanged.

Test plan

  • pytest xrspatial/tests/test_convolution.py (9 passed)
  • New parametrized test asserts the dask declared dtype and the
    computed chunk dtype both equal the eager dtype for int32, float32,
    and float64 input.

)

The dask backends declared float64 via meta=np.array(()) even when the
eager numpy/cupy backends promote an int or float32 raster to float32,
so result.dtype disagreed across backends and did not match the float32
chunks dask actually computes. Pass the promoted float dtype into the
map_overlap meta so the declared dtype matches both the eager backends
and the computed chunks. Adds a parametrized regression test.

Also records the convolution metadata sweep state (2026-07-02).
@github-actions github-actions Bot added the performance PR touches performance-sensitive code label Jul 2, 2026

@brendancol brendancol left a comment

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

PR Review: Align convolution dask output dtype with numpy/cupy (#3618)

Domain-aware review (metadata sweep).

Blockers (must fix before merge)

  • None.

Suggestions (should fix, not blocking)

  • None.

Nits (optional improvements)

  • The regression test in xrspatial/tests/test_convolution.py:12 covers only the numpy and dask+numpy backends. The same untyped-meta bug lived on _convolve_2d_dask_cupy (xrspatial/convolution.py:477) and is fixed the same way, but there's no GPU assertion, so the cupy half of the fix rides on inspection. A @cuda_and_cupy_available dtype check would exercise the cupy and dask+cupy paths on hosts with a GPU. Not blocking, since the CI matrix skips GPU tests anyway.

What looks good

  • The change stays scoped to the reported defect. Both dask backends hoist _promote_float(data.dtype) into a local and pass it as the meta dtype, so the lazy DataArray advertises the same dtype the chunks compute (np.empty_like(data) on the already-promoted input in _convolve_2d_numpy). Nothing else in the two functions moves.
  • Verified on a CUDA host across all four backends: int32/float32 input reports float32 everywhere, and the two dask paths used to advertise float64. float64 input is unchanged. pytest xrspatial/tests/test_convolution.py passes.
  • The wider metadata audit came back clean. convolution_2d keeps name, dims, coords (passthrough coords included), and attrs (crs, res, transform, nodatavals, _FillValue) on every backend. dtype was the only cross-backend divergence.
  • The pre-existing isort drift in convolution.py was left alone (it already fails isort on main and belongs in a style pass), and the second commit only fixes the isort grouping the new test imports introduced.

Checklist

  • Algorithm matches reference/paper (no algorithm change; dtype metadata only)
  • All implemented backends produce consistent results (float32 for int/float32, float64 preserved)
  • NaN handling is correct (unchanged)
  • Edge cases are covered by tests (int32/float32/float64 parametrized)
  • Dask chunk boundaries handled correctly (depth/boundary unchanged)
  • No premature materialization or unnecessary copies
  • Benchmark exists or is not needed (metadata-only fix, no perf change)
  • README feature matrix updated (not applicable)
  • Docstrings present and accurate

@brendancol brendancol left a comment

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Follow-up review (#3622)

Follows up on my earlier review's one nit (no GPU dtype assertion) and the sync with main.

What changed since the last review

  • New commit 26c904e8 adds test_convolve_2d_cupy_dtype_matches_numpy in xrspatial/tests/test_convolution.py:26, a @cuda_and_cupy_available parametrization over int32/float32/float64 that asserts the cupy and dask+cupy outputs declare (and compute) the same promoted dtype the eager numpy backend returns. That closes the gap the nit called out, so the meta=cupy.array((), dtype=fdtype) half of the fix is regression-tested rather than eyeballed.
  • Merged origin/main (19 commits). main had touched convolution.py only in a style commit (#3613); the dtype fix is unchanged and there were no conflicts.

Verification

  • Ran the new GPU test live on a CUDA host (not skipped): all three dtype cases pass on both cupy and dask+cupy. The full xrspatial/tests/test_convolution.py is 12 passed. isort and flake8 are clean on the changed files after the merge.

No remaining findings.

@brendancol brendancol merged commit 59ef9d2 into xarray-contrib:main Jul 6, 2026
10 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

performance PR touches performance-sensitive code

Projects

None yet

Development

Successfully merging this pull request may close these issues.

convolution_2d dask backend reports float64 dtype while numpy/cupy return float32

1 participant