[MPS] Add support for aten::median for MPS backend

[DenisVieriu97]

🐛 Describe the bug

First time contributors are welcome! 🙂
Add support for aten::median for MPS backend. Generic support for adding operations to MPS backend is captured here: https://github.com/pytorch/pytorch/wiki/MPS-Backend#adding-op-for-mps-backend

Versions

N/A

cc @kulinseth @albanD @malfet @razarmehr @abhudev

[Raman-Kumar]

@DenisVieriu97
I want to take and work on this issue.
started working on it ...

[Raman-Kumar]

I sent a Pull request for torch.median(input_tensor). test passed
pytorch/pull/88554

Now I am working on median with dim value
For that I used sortWithTensor and argSortWithTensor from MPSGraph

But it will only work on macOS Ventura

[Raman-Kumar]

will check macOS Ventura like this PR

I need a small help in
the implementation function signature of
torch.median(input, dim=- 1, keepdim=False, *, out=None)

Successfully added median_mps of entire tensor and its implementation function

Now trying to add median op with dim
I added Here MPS: median_out_mps

pytorch/aten/src/ATen/native/native_functions.yaml

Lines 3500 to 3505 in 97be8da

	- func: median.dim_values(Tensor self, int dim, bool keepdim=False, *, Tensor(a!) values, Tensor(b!) indices) -> (Tensor(a!) values, Tensor(b!) indices)
	dispatch:
	CPU: median_out_cpu
	CUDA: median_out_cuda

and gave function signature like below replacing my median op name

pytorch/aten/src/ATen/native/mps/operations/ReduceOps.mm

Lines 1374 to 1386 in 7382c88

	// Max out with dim
	TORCH_IMPL_FUNC(max_out_mps)
	(const Tensor& input_t,
	int64_t dim,
	bool keepdim,
	const Tensor& output_t,
	const Tensor& indices_t) {
	int64_t dim_ = maybe_wrap_dim(dim, input_t.dim());
	native::zero_numel_check_dims(input_t, dim_, "max()");
	min_max_out_mps(input_t, dim, keepdim, output_t, indices_t, MPSReductionType::MAX, "max_out_mps");
	}

But this gives error while building

../aten/src/ATen/native/mps/operations/ReduceOps.mm:1392:1: error: use of undeclared identifier 'structured_median_out_mps'; did you mean 'structured_mean_out_mps'?
TORCH_IMPL_FUNC(median_out_mps)
^
../aten/src/ATen/TensorMeta.h:63:36: note: expanded from macro 'TORCH_IMPL_FUNC'
#define TORCH_IMPL_FUNC(name) void structured_##name::impl
                                   ^
<scratch space>:17:1: note: expanded from here
structured_median_out_mps
^
aten/src/ATen/ops/mean_native.h:23:18: note: 'structured_mean_out_mps' declared here
struct TORCH_API structured_mean_out_mps : public at::meta::structured_mean_dim {
                 ^
../aten/src/ATen/native/mps/operations/ReduceOps.mm:1392:1: error: out-of-line definition of 'impl' does not match any declaration in 'at::native::structured_mean_out_mps'
TORCH_IMPL_FUNC(median_out_mps)
^~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
../aten/src/ATen/TensorMeta.h:63:55: note: expanded from macro 'TORCH_IMPL_FUNC'
#define TORCH_IMPL_FUNC(name) void structured_##name::impl

What should be median_out_mps with dim implementation function signature
and Where it should be placed in /native_functions.yaml?

[Raman-Kumar]

This function signature is working

TORCH_API ::std::tuple<at::Tensor &,at::Tensor &> median_out_mps
(const at::Tensor & self,
 int64_t dim,
bool keepdim, 
at::Tensor & values, 
at::Tensor & indices){

  return std::tuple<Tensor&, Tensor&>(values, indices);
}

Going ahead with this.