Modified median logic to use more suitable

Author: Dhruvanshu-Joshi

pytensor/tensor/math.py

@@ -1588,20 +1588,20 @@ def median(input, axis=None):
     -----
     This function uses the numpy implementation of median.
     """
+    from pytensor.ifelse import ifelse
 
     input = as_tensor_variable(input)
     sorted_input = input.sort(axis=axis)
     shape = input.shape[axis]
     k = extract_constant(shape) // 2
-    if extract_constant(shape % 2) == 0:
-        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)
-        median_val = (ans1 + ans2) / 2.0
-    else:
-        indices = expand_dims(full_like(sorted_input.take(0, axis=axis), k), axis)
-        median_val = take_along_axis(sorted_input, indices, axis=axis)
+    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)
+    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.name = "median"
     return median_val.squeeze(axis=axis)