PaddlePaddle
diff --git a/‎paddle/phi/api/yaml/backward.yaml
Copy file name to clipboard
+2-2Lines changed: 2 additions & 2 deletions b/‎paddle/phi/api/yaml/backward.yaml
Copy file name to clipboard
+2-2Lines changed: 2 additions & 2 deletions
diff --git a/‎paddle/phi/api/yaml/ops.yaml
Copy file name to clipboard
+1-2Lines changed: 1 addition & 2 deletions b/‎paddle/phi/api/yaml/ops.yaml
Copy file name to clipboard
+1-2Lines changed: 1 addition & 2 deletions
diff --git a/‎paddle/phi/infermeta/backward.cc
Copy file name to clipboard
+1Lines changed: 1 addition & 0 deletions b/‎paddle/phi/infermeta/backward.cc
Copy file name to clipboard
+1Lines changed: 1 addition & 0 deletions
diff --git a/‎paddle/phi/infermeta/backward.h
Copy file name to clipboard
+1Lines changed: 1 addition & 0 deletions b/‎paddle/phi/infermeta/backward.h
Copy file name to clipboard
+1Lines changed: 1 addition & 0 deletions
diff --git a/‎paddle/phi/infermeta/unary.cc
Copy file name to clipboard
+9-4Lines changed: 9 additions & 4 deletions b/‎paddle/phi/infermeta/unary.cc
Copy file name to clipboard
+9-4Lines changed: 9 additions & 4 deletions
diff --git a/‎paddle/phi/infermeta/unary.h
Copy file name to clipboard
+1Lines changed: 1 addition & 0 deletions b/‎paddle/phi/infermeta/unary.h
Copy file name to clipboard
+1Lines changed: 1 addition & 0 deletions
diff --git a/‎paddle/phi/kernels/cpu/nanmedian_grad_kernel.cc
Copy file name to clipboard
+46-15Lines changed: 46 additions & 15 deletions b/‎paddle/phi/kernels/cpu/nanmedian_grad_kernel.cc
Copy file name to clipboard
+46-15Lines changed: 46 additions & 15 deletions
diff --git a/‎paddle/phi/kernels/cpu/nanmedian_kernel.cc
Copy file name to clipboard
+58-21Lines changed: 58 additions & 21 deletions b/‎paddle/phi/kernels/cpu/nanmedian_kernel.cc
Copy file name to clipboard
+58-21Lines changed: 58 additions & 21 deletions
@@ -1647,8 +1647,8 @@
     func : mv_grad
 
 - backward_op : nanmedian_grad
-  forward : nanmedian (Tensor x, IntArray axis, bool keepdim) -> Tensor(out), Tensor(medians)
-  args : (Tensor x, Tensor medians, Tensor out_grad, IntArray axis, bool keepdim)
+  forward : nanmedian (Tensor x, IntArray axis, bool keepdim, str mode) -> Tensor(out), Tensor(medians)
+  args : (Tensor x, Tensor medians, Tensor out_grad, IntArray axis, bool keepdim, str mode)
   output : Tensor(x_grad)
   infer_meta :
     func : NanmedianGradInferMeta
 
@@ -2034,13 +2034,12 @@
   backward : mv_grad
 
 - op : nanmedian
-  args : (Tensor x, IntArray axis = {}, bool keepdim = true)
+  args : (Tensor x, IntArray axis = {}, bool keepdim = true, str mode="avg")
   output : Tensor(out), Tensor(medians)
   infer_meta :
     func : NanmedianInferMeta
   kernel :
     func : nanmedian
-  intermediate : medians
   backward : nanmedian_grad
 
 - op : nearest_interp
 
@@ -843,6 +843,7 @@ void NanmedianGradInferMeta(const MetaTensor& x,
                             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);
 
@@ -370,6 +370,7 @@ void NanmedianGradInferMeta(const MetaTensor& x,
                             const MetaTensor& out_grad,
                             const IntArray& axes,
                             bool keep_dim,
+                            const std::string& mode,
                             MetaTensor* x_grad);
 
 void NceGradInferMeta(const MetaTensor& input,
 
@@ -2584,14 +2584,12 @@ void MultinomialInferMeta(const MetaTensor& x,
 void NanmedianInferMeta(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();
-  out->set_dtype(x.dtype());
-  median_index->set_dtype(DataType::INT64);
-  median_index->set_dims(common::make_ddim({x.numel() * 2}));
 
   std::vector<int32_t> out_dim;
   if (axis_list.empty()) {
@@ -2646,8 +2644,15 @@ void NanmedianInferMeta(const MetaTensor& x,
       }
     }
   }
+  out->set_dtype(x.dtype());
+  out->set_dims(make_ddim(out_dim));
 
-  out->set_dims(common::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 NMSInferMeta(const MetaTensor& x, float threshold, MetaTensor* out) {
 
@@ -396,6 +396,7 @@ void MultinomialInferMeta(const MetaTensor& x,
 void NanmedianInferMeta(const MetaTensor& x,
                         const IntArray& axes,
                         bool keep_dim,
+                        const std::string& mode,
                         MetaTensor* out,
                         MetaTensor* median_index);
 
 
@@ -21,11 +21,50 @@
 
 namespace phi {
 
+template <typename T>
+void CalcMedianMeanGrad(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[2 * i] >= 0) {
+      if (m_data[2 * i] == m_data[2 * i + 1]) {
+        dx_data[offset + m_data[2 * i]] = dout_data[i];
+      } else {
+        dx_data[offset + m_data[2 * i]] = dout_data[i] / static_cast<T>(2.0);
+        dx_data[offset + m_data[2 * i + 1]] =
+            dout_data[i] / static_cast<T>(2.0);
+      }
+    }
+    offset += stride;
+  }
+}
+
+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, typename Context>
 void CalcMedianGradKernel(const Context& dev_ctx,
                           const DenseTensor& x,
                           const DenseTensor& median_index,
                           const DenseTensor& out_grad,
+                          const std::string& mode,
                           DenseTensor* x_grad) {
   T* dx_data = dev_ctx.template Alloc<T>(x_grad);
   if (!dx_data) return;
@@ -41,19 +80,10 @@ void CalcMedianGradKernel(const Context& dev_ctx,
   int64_t stride = x_dim[static_cast<int>(rank - 1)];
   int64_t pre_dim = numel / stride;
 
-  int64_t i = 0;
-  int64_t offset = 0;
-  for (i = 0; i < pre_dim; i++) {
-    if (m_data[2 * i] >= 0) {
-      if (m_data[2 * i] == m_data[2 * i + 1]) {
-        dx_data[offset + m_data[2 * i]] = dout_data[i];
-      } else {
-        dx_data[offset + m_data[2 * i]] = dout_data[i] / static_cast<T>(2.0);
-        dx_data[offset + m_data[2 * i + 1]] =
-            dout_data[i] / static_cast<T>(2.0);
-      }
-    }
-    offset += stride;
+  if (mode == "avg") {
+    CalcMedianMeanGrad(pre_dim, stride, m_data, dx_data, dout_data);
+  } else {
+    CalcMedianMinGrad(pre_dim, stride, m_data, dx_data, dout_data);
   }
 }
 
@@ -64,21 +94,22 @@ void NanmedianGradKernel(const Context& dev_ctx,
                          const DenseTensor& out_grad,
                          const IntArray& axes,
                          bool keepdim UNUSED,
+                         const std::string& mode,
                          DenseTensor* x_grad) {
   DenseTensor tmp_x;
   auto rank = x.dims().size();
   if ((axes.size() == 0) || rank <= 1) {
     tmp_x = x;
     tmp_x.Resize({x.numel()});
     CalcMedianGradKernel<T, Context>(
-        dev_ctx, tmp_x, median_index, out_grad, x_grad);
+        dev_ctx, tmp_x, median_index, out_grad, mode, 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<T, Context>(
-        dev_ctx, tmp_x, median_index, out_grad, &tmp_x_grad);
+        dev_ctx, tmp_x, median_index, out_grad, mode, &tmp_x_grad);
 
     dev_ctx.template Alloc<T>(x_grad);
     funcs::PostprocessMedianGradKernel<T, Context>(
 
@@ -30,7 +30,8 @@ void CalcMedianFunc(const Context& dev_ctx,
                     int64_t stride,
                     int64_t pre_dim,
                     T* o_ptr,
-                    int64_t* m_ptr) {
+                    int64_t* m_ptr,
+                    const std::string& mode) {
   DenseTensor sort_out;
   DenseTensor sort_indices;
   auto sort_dim = x.dims();
@@ -51,12 +52,16 @@ void CalcMedianFunc(const Context& dev_ctx,
   int64_t offset = 0;
   int64_t i = 0;
   bool is_ori_odd = stride & 1;
-  if (ignore_nan) {
+  if (ignore_nan) {  // ignore_nan - has nan value; sort_k = max_valid_num
     for (i = 0; i < pre_dim; i++) {
       offset = i * sort_k;
       if (nan_counts[i] == stride) {
-        m_ptr[i * 2] = -1;
-        m_ptr[i * 2 + 1] = -1;
+        if (mode == "avg") {
+          m_ptr[i * 2] = -1;
+          m_ptr[i * 2 + 1] = -1;  // index is -1
+        } else {
+          m_ptr[i] = -1;
+        }
         o_ptr[i] = sort_out_ptr[offset];
       } else {
         int64_t nan_k = nan_counts[i] > 0
@@ -65,21 +70,34 @@ void CalcMedianFunc(const Context& dev_ctx,
         int64_t row_pos = static_cast<int64_t>(nan_k >> 1);
         int64_t pos = offset + row_pos;
         if (nan_k & 1) {
-          m_ptr[2 * i] = sort_indices_ptr[pos];
-          m_ptr[2 * i + 1] = sort_indices_ptr[pos];
+          if (mode == "avg") {
+            m_ptr[2 * i] = sort_indices_ptr[pos];
+            m_ptr[2 * i + 1] = sort_indices_ptr[pos];
+          } else {
+            m_ptr[i] = sort_indices_ptr[pos];
+          }
           o_ptr[i] = sort_out_ptr[pos];
         } else {
-          m_ptr[2 * i] =
-              row_pos > 0 ? sort_indices_ptr[pos - 1] : sort_indices_ptr[pos];
-          m_ptr[2 * i + 1] = sort_indices_ptr[pos];
+          // nan_k is even
           T m_val_left =
               row_pos > 0 ? sort_out_ptr[pos - 1] : sort_out_ptr[pos];
           T m_val_right = sort_out_ptr[pos];
-          o_ptr[i] = (m_val_left + m_val_right) / div_factor;
+          if (mode == "avg") {
+            m_ptr[2 * i] =
+                row_pos > 0 ? sort_indices_ptr[pos - 1] : sort_indices_ptr[pos];
+            m_ptr[2 * i + 1] = sort_indices_ptr[pos];
+            o_ptr[i] = (m_val_left + m_val_right) / div_factor;
+          } else {
+            // mode == "min": output median value should be the left val since
+            // the sort_out is in ascending order
+            m_ptr[i] =
+                row_pos > 0 ? sort_indices_ptr[pos - 1] : sort_indices_ptr[pos];
+            o_ptr[i] = m_val_left;
+          }
         }
       }
     }
-  } else {
+  } else {  // not ignore_nan - no nan value; sort_k = stride/2 + 1
     if (is_ori_odd) {
       for (i = 0; i < pre_dim; i++) {
         offset = i * sort_k;
@@ -92,12 +110,20 @@ void CalcMedianFunc(const Context& dev_ctx,
       for (i = 0; i < pre_dim; i++) {
         offset = i * sort_k;
         int64_t pos = offset + sort_k - 1;
-        m_ptr[2 * i] =
-            sort_k > 1 ? sort_indices_ptr[pos - 1] : sort_indices_ptr[pos];
-        m_ptr[2 * i + 1] = sort_indices_ptr[pos];
         T m_val_left = sort_k > 1 ? sort_out_ptr[pos - 1] : sort_out_ptr[pos];
         T m_val_right = sort_out_ptr[pos];
-        o_ptr[i] = (m_val_left + m_val_right) / div_factor;
+        if (mode == "avg") {
+          m_ptr[2 * i] =
+              sort_k > 1 ? sort_indices_ptr[pos - 1] : sort_indices_ptr[pos];
+          m_ptr[2 * i + 1] = sort_indices_ptr[pos];
+          o_ptr[i] = (m_val_left + m_val_right) / div_factor;
+        } else {
+          // mode == "min": output median value should be the left val since the
+          // sort_out is in ascending order
+          m_ptr[i] =
+              sort_k > 1 ? sort_indices_ptr[pos - 1] : sort_indices_ptr[pos];
+          o_ptr[i] = m_val_left;
+        }
       }
     }
   }
@@ -106,6 +132,7 @@ void CalcMedianFunc(const Context& dev_ctx,
 template <typename T, typename Context>
 void ProcessMedianKernel(const Context& dev_ctx,
                          const DenseTensor& x,
+                         const std::string& mode,
                          DenseTensor* out,
                          DenseTensor* median_index) {
   const T* x_data = x.data<T>();
@@ -154,8 +181,12 @@ void ProcessMedianKernel(const Context& dev_ctx,
     if (total_nan_num == numel) {
       for (i = 0; i < pre_dim; i++) {
         out_data[i] = std::numeric_limits<T>::quiet_NaN();
-        m_data[2 * i] = -1;
-        m_data[2 * i + 1] = -1;
+        if (mode == "avg") {
+          m_data[2 * i] = -1;
+          m_data[2 * i + 1] = -1;  // indices are all -1
+        } else {
+          m_data[i] = -1;
+        }
       }
       return;
     }
@@ -171,26 +202,32 @@ void ProcessMedianKernel(const Context& dev_ctx,
                              stride,
                              pre_dim,
                              out_data,
-                             m_data);
+                             m_data,
+                             mode);
 }
 
 template <typename T, typename Context>
 void NanmedianKernel(const Context& dev_ctx,
                      const DenseTensor& x,
                      const IntArray& axes,
                      bool keepdim UNUSED,
+                     const std::string& mode,
                      DenseTensor* out,
                      DenseTensor* median_index) {
   DenseTensor tmp_x;
   auto rank = x.dims().size();
   if ((axes.size() == 0) || rank <= 1) {
     tmp_x = x;
-    tmp_x.Resize({x.numel()});
+    tmp_x.Resize({x.numel()});  // flatten
   } else {
-    funcs::PreprocessMedianKernel<T, Context>(dev_ctx, x, axes, &tmp_x);
+    funcs::PreprocessMedianKernel<T, Context>(
+        dev_ctx,
+        x,
+        axes,
+        &tmp_x);  // resize to 2D so as to compute median on last axis
   }
 
-  ProcessMedianKernel<T, Context>(dev_ctx, tmp_x, out, median_index);
+  ProcessMedianKernel<T, Context>(dev_ctx, tmp_x, mode, out, median_index);
 }
 
 }  // namespace phi