Modify test for median

Author: Dhruvanshu-Joshi

pytensor/tensor/math.py

@@ -26,7 +26,6 @@
     concatenate,
     constant,
     expand_dims,
-    extract_constant,
     full_like,
     stack,
     switch,
@@ -1580,34 +1579,48 @@ def median(input, axis=None):
 
     Parameters
     ----------
+    input: TensorVariable
+        The input tensor.
     axis: None or int or (list of int) (see `Sum`)
         Compute the median along this axis of the tensor.
         None means all axes (like numpy).
-
-    Notes
-    -----
-    This function uses the numpy implementation of median.
     """
     from pytensor.ifelse import ifelse
 
+    x = as_tensor_variable(input)
+    x_ndim = x.type.ndim
     if axis is None:
-        input = input.flatten()
-        axis = 0
+        axis = list(range(x_ndim))
+    else:
+        axis = list(normalize_axis_tuple(axis, x_ndim))
 
-    input = as_tensor_variable(input)
-    sorted_input = input.sort(axis=axis)
-    shape = input.shape[axis]
-    k = extract_constant(shape) // 2
-    indices1 = expand_dims(full_like(sorted_input.take(0, axis=axis), k - 1), axis)
-    indices2 = expand_dims(full_like(sorted_input.take(0, axis=axis), k), axis)
-    ans1 = take_along_axis(sorted_input, indices1, axis=axis)
-    ans2 = take_along_axis(sorted_input, indices2, axis=axis)
+    new_axes_order = [i for i in range(x.ndim) if i not in axis] + list(axis)
+    x = x.dimshuffle(new_axes_order)
+    x_shape = x.shape
+
+    remaining_axis_size = shape(x)[: x.ndim - len(axis)]
+
+    x = x.reshape((*remaining_axis_size, -1))
+
+    # Sort the input tensor along the specified axis
+    sorted_x = x.sort(axis=-1)
+    x_shape = x.shape[-1]
+    k = x_shape // 2
+
+    indices1 = expand_dims(full_like(sorted_x.take(0, axis=-1), k), -1)
+    indices2 = expand_dims(full_like(sorted_x.take(0, axis=-1), k - 1), -1)
+    ans1 = take_along_axis(sorted_x, indices1, axis=-1)
+    ans2 = take_along_axis(sorted_x, indices2, axis=-1)
     median_val_even = (ans1 + ans2) / 2.0
-    indices = expand_dims(full_like(sorted_input.take(0, axis=axis), k), axis)
-    median_val_odd = take_along_axis(sorted_input, indices, axis=axis)
-    median_val = ifelse(eq(mod(shape, 2), 0), median_val_even, median_val_odd)
+
+    median_val_odd = (
+        take_along_axis(sorted_x, indices1, axis=-1) / 1.0
+    )  # Divide by one so that the two dtypes passed in ifelse are compatible
+
+    median_val = ifelse(eq(mod(x_shape, 2), 0), median_val_even, median_val_odd)
     median_val.name = "median"
-    return median_val.squeeze(axis=axis)
+
+    return median_val.squeeze(axis=-1)
 
 
 @scalar_elemwise(symbolname="scalar_maximum")

tests/tensor/test_math.py

@@ -3735,32 +3735,30 @@ def test_nan_to_num(nan, posinf, neginf):
 
 
 @pytest.mark.parametrize(
-    "data, axis",
+    "ndim, axis",
     [
-        # 1D array
-        ([1, 7, 3, 6, 5, 2, 4], None),
-        ([1, 7, 3, 6, 5, 2, 4], 0),
-        # 2D array
-        ([[6, 2], [4, 3], [1, 5]], 0),
-        ([[6, 2], [4, 3], [1, 5]], 1),
-        # 3D array
-        ([[[6, 2, 3], [1, 5, 8], [4, 7, 9]], [[5, 3, 4], [8, 6, 2], [7, 1, 9]]], None),
-        ([[[6, 2, 3], [1, 5, 8], [4, 7, 9]], [[5, 3, 4], [8, 6, 2], [7, 1, 9]]], 0),
-        ([[[6, 2, 3], [1, 5, 8], [4, 7, 9]], [[5, 3, 4], [8, 6, 2], [7, 1, 9]]], 1),
-        ([[[6, 2, 3], [1, 5, 8], [4, 7, 9]], [[5, 3, 4], [8, 6, 2], [7, 1, 9]]], 2),
-        # 4D array
-        (
-            [
-                [[[3, 1], [4, 3]], [[0, 5], [6, 2]], [[7, 8], [9, 4]]],
-                [[[10, 11], [12, 13]], [[14, 15], [16, 17]], [[18, 19], [20, 21]]],
-            ],
-            3,
-        ),
+        (2, None),
+        (2, 1),
+        (2, (0, 1)),
+        (3, None),
+        (3, (1, 2)),
+        (4, (1, 3, 0)),
     ],
 )
-def test_median(data, axis):
-    x = tensor(shape=np.array(data).shape)
+def test_median(ndim, axis):
+    # Generate random data with both odd and even lengths
+    shape_even = np.arange(1, ndim + 1) * 2
+    shape_odd = shape_even - 1
+
+    data_even = np.random.rand(*shape_even)
+    data_odd = np.random.rand(*shape_odd)
+
+    x = tensor(dtype="float64", shape=(None,) * ndim)
     f = function([x], median(x, axis=axis))
-    result = f(data)
-    expected = np.median(data, axis=axis)
-    assert np.allclose(result, expected)
+    result_odd = f(data_odd)
+    result_even = f(data_even)
+    expected_odd = np.median(data_odd, axis=axis)
+    expected_even = np.median(data_even, axis=axis)
+
+    assert np.allclose(result_odd, expected_odd)
+    assert np.allclose(result_even, expected_even)