Repository: ivy-llc/ivy

Issue:
median
The PyTorch front end median function implemented in #13212 does not follow the correct behaviour. PyTorch's median returns a named tuple with the median and their indices when a axis is provided to reduce along. The incorrect implementation has been reverted in #13480 

https://pytorch.org/docs/stable/generated/torch.median.html


Discussion:

Resolving patch:
diff --git a/ivy/functional/frontends/torch/reduction_ops.py b/ivy/functional/frontends/torch/reduction_ops.py
--- a/ivy/functional/frontends/torch/reduction_ops.py
+++ b/ivy/functional/frontends/torch/reduction_ops.py
@@ -62,6 +62,34 @@ def nanmean(input, dim=None, keepdim=False, *, dtype=None, out=None):
     return ivy.nanmean(input, axis=dim, keepdims=keepdim, dtype=dtype, out=out)
 
 
+@to_ivy_arrays_and_back
+def median(input, dim=None, keepdim=False, *, out=None):
+    if dim is None:
+        input = ivy.reshape(input, (-1,))
+        sorted_input = ivy.sort(input)
+        return sorted_input[(sorted_input.shape[0] - 1) // 2]
+
+    median_tuple = namedtuple("median", ["values", "indices"])
+
+    if input.ndim == 0:
+        result = median_tuple(input, ivy.array(0))
+    else:
+        sorted_indices = ivy.argsort(input, axis=dim)
+        median_indices = ivy.gather(sorted_indices, (sorted_indices.shape[dim] - 1) // 2, axis=dim)
+        median_values = ivy.take_along_axis(input, ivy.expand_dims(median_indices, axis=dim), dim).squeeze(dim)
+
+        if keepdim:
+            median_values = ivy.expand_dims(median_values, axis=dim)
+            median_indices = ivy.expand_dims(median_indices, axis=dim)
+
+        result = median_tuple(median_values, median_indices)
+    if out is not None:
+        ivy.inplace_update(out[0], result.values)
+        ivy.inplace_update(out[1], result.indices)
+        return out
+    return result
+
+
 @to_ivy_arrays_and_back
 def std(input, dim, unbiased, keepdim=False, *, out=None):
     return ivy.std(input, axis=dim, correction=int(unbiased), keepdims=keepdim, out=out)


Test patch:
diff --git a/ivy_tests/test_ivy/test_frontends/test_torch/test_reduction_ops.py b/ivy_tests/test_ivy/test_frontends/test_torch/test_reduction_ops.py
--- a/ivy_tests/test_ivy/test_frontends/test_torch/test_reduction_ops.py
+++ b/ivy_tests/test_ivy/test_frontends/test_torch/test_reduction_ops.py
@@ -335,6 +335,38 @@ def test_torch_nanmean(
     )
 
 
+@handle_frontend_test(
+    fn_tree="torch.median",
+    dtype_input_axis=helpers.dtype_values_axis(
+        available_dtypes=helpers.get_dtypes("numeric"),
+        min_num_dims=1,
+        valid_axis=True,
+        force_int_axis=True,
+    ),
+    keepdim=st.booleans(),
+)
+def test_torch_median(
+    *,
+    dtype_input_axis,
+    keepdim,
+    on_device,
+    fn_tree,
+    frontend,
+    test_flags,
+):
+    input_dtype, input, dim = dtype_input_axis
+    helpers.test_frontend_function(
+        input_dtypes=input_dtype,
+        frontend=frontend,
+        test_flags=test_flags,
+        fn_tree=fn_tree,
+        on_device=on_device,
+        input=input[0],
+        dim=dim,
+        keepdim=keepdim,
+    )
+
+
 @handle_frontend_test(
     fn_tree="torch.std",
     dtype_and_x=statistical_dtype_values(function="std"),