Skip to content

hypsometric_integral dask+cupy backend silently returns dask+numpy result #2525

@brendancol

Description

@brendancol

Description

_hi_dask_cupy converts cupy chunks to numpy via map_blocks(.get()) and then calls _hi_dask_numpy. The returned dask graph yields numpy chunks, so the wrapping xr.DataArray is dask+numpy backed even though the caller passed dask+cupy.

is_dask_cupy(result) returns False. Downstream code that dispatches on backend type (another zonal function using ArrayTypeFunctionMapping) will route through the numpy path. A GPU-resident pipeline loses GPU residency at this hop and the caller never finds out.

_regions_dask_cupy in the same module does this correctly: it materializes on the GPU and returns a dask-cupy result.

Reproducer

import numpy as np
import cupy as cp
import dask.array as da
import xarray as xr
from xrspatial.zonal import hypsometric_integral
from xrspatial.utils import is_dask_cupy

zones_np = np.array([[1, 1, 2, 2], [1, 1, 2, 2]], dtype=np.int32)
values_np = np.array([[10.0, 20.0, 30.0, 40.0], [15.0, 25.0, 35.0, 45.0]])

zones = xr.DataArray(
    da.from_array(cp.asarray(zones_np), chunks=(2, 2)), dims=['y', 'x'],
)
values = xr.DataArray(
    da.from_array(cp.asarray(values_np), chunks=(2, 2)), dims=['y', 'x'],
)

assert is_dask_cupy(zones)
assert is_dask_cupy(values)

result = hypsometric_integral(zones, values)
assert is_dask_cupy(result)  # FAILS: result is dask+numpy

Expected

A dask+cupy input should produce a dask+cupy output. is_dask_cupy(result) should be True and result.data.blocks[0, 0].compute() should return a cupy.ndarray.

Suggested fix

After _hi_dask_cupy calls _hi_dask_numpy, wrap the result through map_blocks(cupy.asarray, ...) so every chunk lands back as a cupy array. The dask graph stays lazy. The numpy reduce in the middle stays as-is (the per-zone HI lookup is a Python dict), but the painted output matches the input backend.

Audit trail

Surfaced by /sweep-metadata on the zonal module, 2026-05-27. Cat 5 (backend-inconsistent metadata), severity MEDIUM.

Metadata

Metadata

Assignees

No one assigned

    Labels

    backend-coverageAdding missing dask/cupy/dask+cupy backend supportbugSomething isn't workinggpuCuPy / CUDA GPU supportzonal tools

    Type

    No type
    No fields configured for issues without a type.

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions