为median添加内核

paddle/phi/infermeta/backward.cc

@@ -1135,6 +1135,19 @@ void MaxPoolWithIndexGradInferMeta(const MetaTensor& x,
   dx->share_meta(x);
 }
 
+void MedianGradInferMeta(const MetaTensor& x,
+                         const MetaTensor& median_data,
+                         const MetaTensor& median_index,
+                         const MetaTensor& out_grad,
+                         const IntArray& axes,
+                         bool keep_dim,
+                         const std::string& mode,
+                         MetaTensor* x_grad) {
+  auto x_dims = x.dims();
+  x_grad->set_dims(x_dims);
+  x_grad->set_dtype(x.dtype());
+}
+
 void MemoryEfficientAttentionGradInferMeta(const MetaTensor& query,
                                            const MetaTensor& key,
                                            const MetaTensor& value,
@@ -1417,6 +1430,7 @@ void MultiplexGradInferMeta(const MetaTensor& ids,
 }
 
 void NanmedianGradInferMeta(const MetaTensor& x,
+                            const MetaTensor& median_data,
                             const MetaTensor& median_index,
                             const MetaTensor& out_grad,
                             const IntArray& axes,

paddle/phi/infermeta/backward.h

@@ -446,6 +446,15 @@ void MaxPoolWithIndexGradInferMeta(const MetaTensor& x,
                                    bool ceil_mode,
                                    MetaTensor* dx);
 
+void MedianGradInferMeta(const MetaTensor& x,
+                         const MetaTensor& median_data,
+                         const MetaTensor& median_index,
+                         const MetaTensor& out_grad,
+                         const IntArray& axes,
+                         bool keep_dim,
+                         const std::string& mode,
+                         MetaTensor* x_grad);
+
 void MeshgridGradInferMeta(const std::vector<const MetaTensor*>& inputs,
                            const std::vector<const MetaTensor*>& outputs_grad,
                            std::vector<MetaTensor*> inputs_grad);
@@ -525,6 +534,7 @@ void MultiplexGradInferMeta(const MetaTensor& ids,
                             std::vector<MetaTensor*> ins_grad);
 
 void NanmedianGradInferMeta(const MetaTensor& x,
+                            const MetaTensor& median_data,
                             const MetaTensor& median_index,
                             const MetaTensor& out_grad,
                             const IntArray& axes,

paddle/phi/infermeta/unary.cc

@@ -2890,6 +2890,80 @@ void MeanAllInferMeta(const MetaTensor& x, MetaTensor* out) {
   out->set_layout(x.layout());
 }
 
+void MedianInferMeta(const MetaTensor& x,
+                     const IntArray& axes,
+                     bool keep_dim,
+                     const std::string& mode,
+                     MetaTensor* out,
+                     MetaTensor* median_index) {
+  std::vector<int64_t> axis_list = axes.GetData();
+  auto x_dim = x.dims();
+  int64_t x_rank = x_dim.size();
+
+  std::vector<int64_t> out_dim;
+  if (axis_list.empty()) {
+    if (keep_dim) {
+      for (int64_t i = 0; i < x_rank; i++) {
+        out_dim.push_back(1);
+      }
+    }
+  } else {
+    std::vector<int64_t> formatted_axis;
+    for (auto& axis : axis_list) {
+      if (x_rank == 0) {
+        PADDLE_ENFORCE_EQ(axis == 0 || axis == -1,
+                          true,
+                          common::errors::InvalidArgument(
+                              "When input 0D Tensor, each element of the axis "
+                              "can only be -1, 0, None"));
+      } else {
+        PADDLE_ENFORCE_LT(axis,
+                          x_rank,
+                          errors::InvalidArgument(
+                              "each element of the axis should be in the "
+                              "range [ -dimension(X), dimension(X) ) "
+                              "which dimension = %d. But received axis = %d.",
+                              x_rank,
+                              axis));
+        PADDLE_ENFORCE_GE(axis,
+                          -x_rank,
+                          errors::InvalidArgument(
+                              "each element of the axis should be in the "
+                              "range [ -dimension(X), dimension(X) ) "
+                              "which dimension = %d. But received axis = %d.",
+                              x_rank,
+                              axis));
+      }
+      if (axis < 0) axis += x_rank;
+      PADDLE_ENFORCE_EQ(
+          std::find(formatted_axis.begin(), formatted_axis.end(), axis),
+          formatted_axis.end(),
+          errors::InvalidArgument("Attr(axes) has duplicated elements: %d.",
+                                  static_cast<int>(axis)));
+
+      formatted_axis.push_back(axis);
+    }
+
+    for (int64_t i = 0; i < x_rank; i++) {
+      if (std::find(formatted_axis.begin(), formatted_axis.end(), i) ==
+          formatted_axis.end()) {
+        out_dim.push_back(x_dim[i]);  // NOLINT
+      } else if (keep_dim) {
+        out_dim.push_back(1);
+      }
+    }
+  }
+  out->set_dtype(x.dtype());
+  out->set_dims(make_ddim(out_dim));
+
+  auto median_dim = out_dim;
+  if (mode == "avg") {
+    median_dim.push_back(2);
+  }
+  median_index->set_dtype(DataType::INT64);
+  median_index->set_dims(make_ddim(median_dim));
+}
+
 void ModeInferMeta(const MetaTensor& x,
                    int axis,
                    bool keepdim,

paddle/phi/infermeta/unary.h

@@ -468,6 +468,13 @@ void MaxPoolV2InferMeta(const MetaTensor& x,
 
 void MeanAllInferMeta(const MetaTensor& x, MetaTensor* out);
 
+void MedianInferMeta(const MetaTensor& x,
+                     const IntArray& axes,
+                     bool keep_dim,
+                     const std::string& mode,
+                     MetaTensor* out,
+                     MetaTensor* median_index);
+
 void ModeInferMeta(const MetaTensor& x,
                    int axis,
                    bool keepdim,

paddle/phi/kernels/cpu/median_grad_kernel.cc

@@ -0,0 +1,169 @@
+// Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "paddle/phi/kernels/median_grad_kernel.h"
+
+#include <math.h>
+#include "paddle/phi/backends/cpu/cpu_context.h"
+#include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/kernels/funcs/math_function.h"
+#include "paddle/phi/kernels/funcs/nanmedian_utils.h"
+
+namespace phi {
+
+template <typename T>
+void CalcMedianMinGrad(int64_t pre_dim,
+                       int64_t stride,
+                       const int64_t* m_data,
+                       T* dx_data,
+                       const T* dout_data) {
+  int64_t i = 0;
+  int64_t offset = 0;
+  for (i = 0; i < pre_dim; i++) {
+    if (m_data[i] >= 0) {
+      dx_data[offset + m_data[i]] = dout_data[i];
+    }
+    offset += stride;
+  }
+}
+
+template <typename T>
+void CalcMedianGradEvenly(int64_t pre_dim,
+                          int64_t stride,
+                          const DenseTensor& x,
+                          const T* m_data,
+                          const int64_t* m_index,
+                          T* dx_data,
+                          const T* dout_data) {
+  int64_t i = 0, j = 0;
+  int64_t offset = 0;
+  std::vector<int64_t> data_index;
+  const T* x_data = x.data<T>();
+  for (i = 0; i < pre_dim; i++) {
+    data_index.clear();
+    for (j = 0; j < stride; j++) {
+      if ((m_data[i] == x_data[offset + j]) ||
+          (isnan(static_cast<float>(m_data[i])) &&
+           isnan(static_cast<float>(x_data[offset + j])))) {
+        data_index.push_back(offset + j);
+      }
+    }
+    if (data_index.size() == 0) {
+      if (m_index[2 * i] == m_index[2 * i + 1]) {
+        dx_data[offset + m_index[2 * i]] = dout_data[i];
+      } else {
+        dx_data[offset + m_index[2 * i]] = dout_data[i] / static_cast<T>(2.0);
+        dx_data[offset + m_index[2 * i + 1]] =
+            dout_data[i] / static_cast<T>(2.0);
+      }
+    } else {
+      for (j = 0; j < data_index.size(); j++) {
+        dx_data[data_index[j]] =
+            dout_data[i] / static_cast<T>(data_index.size());
+      }
+    }
+
+    offset += stride;
+  }
+}
+
+template <typename T, typename Context>
+void CalcMedianGradKernel_CPU(const Context& dev_ctx,
+                              const DenseTensor& x,
+                              const DenseTensor& median_data,
+                              const DenseTensor& median_index,
+                              const DenseTensor& out_grad,
+                              const std::string& mode,
+                              const bool evenly,
+                              DenseTensor* x_grad) {
+  T* dx_data = dev_ctx.template Alloc<T>(x_grad);
+  if (!dx_data) return;
+
+  phi::funcs::SetConstant<Context, T> set_zero;
+  set_zero(dev_ctx, x_grad, static_cast<T>(0));
+
+  const int64_t* m_index = median_index.data<int64_t>();
+  const T* m_data = median_data.data<T>();
+  const T* dout_data = out_grad.data<T>();
+  int64_t numel = x.numel();
+  auto x_dim = x.dims();
+  int64_t rank = x_dim.size();
+  int64_t stride = x_dim[static_cast<int>(rank - 1)];
+  int64_t pre_dim = numel / stride;
+  if (!evenly) {
+    CalcMedianMinGrad(pre_dim, stride, m_index, dx_data, dout_data);
+  } else {
+    CalcMedianGradEvenly(
+        pre_dim, stride, x, m_data, m_index, dx_data, dout_data);
+  }
+}
+
+template <typename T, typename Context>
+void MedianGradKernel(const Context& dev_ctx,
+                      const DenseTensor& x,
+                      const DenseTensor& median_data,
+                      const DenseTensor& median_index,
+                      const DenseTensor& out_grad,
+                      const IntArray& axes,
+                      bool keepdim UNUSED,
+                      const std::string& mode,
+                      DenseTensor* x_grad) {
+  if (x_grad && x_grad->numel() == 0) {
+    dev_ctx.template Alloc<T>(x_grad);
+    return;
+  }
+  bool evenly = (axes.size() != 1 || mode == "avg");
+  DenseTensor tmp_x;
+  auto rank = x.dims().size();
+  if ((axes.size() == 0) || rank <= 1) {
+    tmp_x = x;
+    tmp_x.Resize({x.numel()});
+    CalcMedianGradKernel_CPU<T, Context>(dev_ctx,
+                                         tmp_x,
+                                         median_data,
+                                         median_index,
+                                         out_grad,
+                                         mode,
+                                         evenly,
+                                         x_grad);
+  } else {
+    funcs::PreprocessMedianKernel<T, Context>(dev_ctx, x, axes, &tmp_x);
+
+    DenseTensor tmp_x_grad;
+    tmp_x_grad.Resize(x_grad->dims());
+    CalcMedianGradKernel_CPU<T, Context>(dev_ctx,
+                                         tmp_x,
+                                         median_data,
+                                         median_index,
+                                         out_grad,
+                                         mode,
+                                         evenly,
+                                         &tmp_x_grad);
+
+    dev_ctx.template Alloc<T>(x_grad);
+    funcs::PostprocessMedianGradKernel<T, Context>(
+        dev_ctx, &tmp_x_grad, axes, x_grad);
+  }
+}
+
+}  // namespace phi
+
+PD_REGISTER_KERNEL(median_grad,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::MedianGradKernel,
+                   float,
+                   double,
+                   int,
+                   int64_t) {}