Support None

Author: Dhruvanshu-Joshi

pytensor/tensor/math.py

@@ -1590,6 +1590,10 @@ def median(input, axis=None):
     """
     from pytensor.ifelse import ifelse
 
+    if axis is None:
+        input = input.flatten()
+        axis = 0
+
     input = as_tensor_variable(input)
     sorted_input = input.sort(axis=axis)
     shape = input.shape[axis]

tests/tensor/test_math.py

@@ -3738,11 +3738,13 @@ def test_nan_to_num(nan, posinf, neginf):
     "data, 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),