Fix dpnp.median/dpnp.nanmedian result shape for empty and size-1 kept dimensions (#3081)

This PR fixes two kept-dimension shape bugs in
`dpnp.median`/`dpnp.nanmedian`.

**Empty kept dimension.** `dpnp.median` and `dpnp.nanmedian` raised
`ValueError` when called with a tuple (or list) `axis` and one of the
kept, non-reduced dimensions has size 0, e.g.
```
dpnp.median(dpnp.empty((0, 3, 4)), axis=(1, 2))
```
The sequence-axis path flattens the reduced axes in
`_flatten_array_along_axes` with `a.reshape(kept_shape + (-1,))`. The
`-1` cannot be inferred once the array has size 0 and a kept axis is
also 0 (every merged length yields a size-0 array, so the dimension is
ambiguous), which makes `reshape` raise. The merged length is now
computed explicitly with `math.prod(...)`, so the tuple-axis path
returns the same empty result as the single-axis path.

**Size-1 kept dimension.** `dpnp.nanmedian` dropped kept dimensions of
size 1 and returned the wrong shape, e.g. a `(1, 5)` array reduced over
`axis=1` returned shape `()` instead of `(1,)`. `_calc_nanmedian` ended
with an unconditional `dpnp.squeeze(res)`, which also removed
non-reduced size-1 axes; it now squeezes only the reduced trailing axis
with `dpnp.squeeze(res, axis=-1)`.

Author: antonwolfy

CHANGELOG.md

@@ -110,7 +110,9 @@ This release is compatible with NumPy 2.5.
 * Fixed the list of events the copy kernels of `dpnp.reshape`, `dpnp.tensor.reshape`, `dpnp.roll` and `dpnp.tensor.roll` wait on being padded with default-constructed events [#3072](https://github.com/IntelPython/dpnp/pull/3072)
 * Fixed `simplify_iteration_three_strides` and `simplify_iteration_four_strides` accumulating into their third and fourth output displacements without zeroing them first, which required the caller to initialize them [#3072](https://github.com/IntelPython/dpnp/pull/3072)
 * Fixed `dpnp.ndarray.flat` indexing and assignment edge cases, adding support for slices, ellipsis, and integer/boolean array indices [#3045](https://github.com/IntelPython/dpnp/pull/3045)
+* Fixed `dpnp.nanmedian` dropping kept dimensions of size 1, which produced a wrong result shape [#3081](https://github.com/IntelPython/dpnp/pull/3081)
 * Fixed incorrect results of `dpnp.tensor.vecdot` in some cases with strided outputs and of `dpnp.tensor` reductions, `dpnp.tensor.vecdot` and `dpnp.tensor.matmul` on large inputs with some data types [#3082](https://github.com/IntelPython/dpnp/pull/3082)
+* Fixed `dpnp.median` and `dpnp.nanmedian` raising a `ValueError` for a tuple `axis` when a kept dimension has size 0 [#3081](https://github.com/IntelPython/dpnp/pull/3081)
 
 ### Security
 

dpnp/dpnp_utils/dpnp_utils_statistics.py

@@ -26,6 +26,7 @@
 # THE POSSIBILITY OF SUCH DAMAGE.
 # *****************************************************************************
 
+import math
 import warnings
 
 import dpnp
@@ -88,7 +89,8 @@ def _calc_nanmedian(a, out=None):
     if mask.all(axis=-1).any():
         warnings.warn("All-NaN slice encountered", RuntimeWarning, stacklevel=6)
 
-    return dpnp.squeeze(res)
+    # only drop the reduced axis, keep size-1 dimensions that are not reduced
+    return dpnp.squeeze(res, axis=-1)
 
 
 def _flatten_array_along_axes(a, axes_to_flatten, overwrite_input):
@@ -102,7 +104,10 @@ def _flatten_array_along_axes(a, axes_to_flatten, overwrite_input):
     # Move the axes_to_flatten to the end
     destination = list(range(len(axes_to_keep), a_ndim))
     a_moved = dpnp.moveaxis(a, axes_to_flatten, destination)
-    new_shape = tuple(a.shape[axis] for axis in axes_to_keep) + (-1,)
+    # Compute the merged length explicitly instead of letting `reshape` infer
+    # it with -1, since -1 is ambiguous when a kept axis has size 0
+    merged = math.prod(a.shape[axis] for axis in axes_to_flatten)
+    new_shape = tuple(a.shape[axis] for axis in axes_to_keep) + (merged,)
     a_flatten = a_moved.reshape(new_shape)
 
     # Note that the output of a_flatten is not necessarily a view of the input

dpnp/tests/test_nanfunctions.py

@@ -413,6 +413,34 @@ def test_empty(self, axis, shape):
         expected = numpy.nanmedian(a, axis=axis)
         assert_dtype_allclose(result, expected)
 
+    @pytest.mark.usefixtures("suppress_mean_empty_slice_numpy_warnings")
+    @pytest.mark.parametrize(
+        "keepdims, out_shape", [(False, (0,)), (True, (0, 1, 1))]
+    )
+    def test_empty_kept_dim(self, keepdims, out_shape):
+        a = numpy.empty((0, 3, 4))
+        ia = dpnp.array(a)
+
+        result = dpnp.nanmedian(ia, axis=(1, 2), keepdims=keepdims)
+        assert result.shape == out_shape
+        if numpy_version() >= "2.5.4":
+            expected = numpy.nanmedian(a, axis=(1, 2), keepdims=keepdims)
+            assert_dtype_allclose(result, expected)
+
+    @pytest.mark.usefixtures("suppress_mean_empty_slice_numpy_warnings")
+    @pytest.mark.parametrize(
+        "shape, axis", [((1, 5), 1), ((3, 1, 4), 2), ((2, 1, 5), (0, 2))]
+    )
+    @pytest.mark.parametrize("keepdims", [True, False])
+    def test_size1_kept_dim(self, shape, axis, keepdims):
+        a = generate_random_numpy_array(shape)
+        a.flat[0] = numpy.nan
+        ia = dpnp.array(a)
+
+        result = dpnp.nanmedian(ia, axis=axis, keepdims=keepdims)
+        expected = numpy.nanmedian(a, axis=axis, keepdims=keepdims)
+        assert_dtype_allclose(result, expected)
+
     @pytest.mark.parametrize("dtype", get_all_dtypes(no_none=True))
     @pytest.mark.parametrize("axis", [None, 0, (-1,), [0, 1], (0, -2, -1)])
     def test_no_nan(self, dtype, axis):

dpnp/tests/test_statistics.py

@@ -944,6 +944,20 @@ def test_empty(self, axis, shape):
         expected = numpy.median(a, axis=axis)
         assert_dtype_allclose(result, expected)
 
+    @pytest.mark.usefixtures("suppress_mean_empty_slice_numpy_warnings")
+    @pytest.mark.parametrize(
+        "keepdims, out_shape", [(False, (0,)), (True, (0, 1, 1))]
+    )
+    def test_empty_kept_dim(self, keepdims, out_shape):
+        a = numpy.empty((0, 3, 4))
+        ia = dpnp.array(a)
+
+        result = dpnp.median(ia, axis=(1, 2), keepdims=keepdims)
+        assert result.shape == out_shape
+        if numpy_version() >= "2.5.4":
+            expected = numpy.median(a, axis=(1, 2), keepdims=keepdims)
+            assert_dtype_allclose(result, expected)
+
     @pytest.mark.parametrize("dtype", get_all_dtypes())
     @pytest.mark.parametrize(
         "axis, out_shape", [(0, (3,)), (1, (2,)), ((0, 1), ())]