[cherry-pick 2.5][Zero-Dim] paddle.nanmedian/count_nonzero/logspace support...

[cherry-pick 2.5][Zero-Dim] paddle.nanmedian/count_nonzero/logspace support 0D, add some 0D case (#54649)

* [Zero-Dim] add 0D test case (#54581)

* [Zero-Dim] paddle.nanmedian/nanquantile support 0D Tensor (#54500)

* [Zero-Dim] paddle.nanmedian support 0D Tensor

* fix CI

paddle/phi/infermeta/multiary.cc

@@ -2162,32 +2162,32 @@ void LogspaceInferMeta(const MetaTensor& start,
                        MetaTensor* out) {
   auto s_dims = start.dims();
   PADDLE_ENFORCE_EQ(
-      (s_dims.size() == 1) && (s_dims[0] == 1),
-      true,
-      phi::errors::InvalidArgument("The shape of Input(Start) must be [1],"
-                                   "but received input shape is [%s].",
-                                   s_dims));
+      phi::product(s_dims),
+      1,
+      phi::errors::InvalidArgument("The size of Input(Start) must be 1,"
+                                   "but received input size is %s.",
+                                   phi::product(s_dims)));
   auto e_dims = stop.dims();
   PADDLE_ENFORCE_EQ(
-      (e_dims.size() == 1) && (e_dims[0] == 1),
+      phi::product(e_dims),
       true,
-      phi::errors::InvalidArgument("The shape of Input(Stop) must be [1],"
-                                   "but received input shape is [%s].",
-                                   e_dims));
+      phi::errors::InvalidArgument("The size of Input(Stop) must be 1,"
+                                   "but received input size is %s.",
+                                   phi::product(e_dims)));
   auto num_dims = number.dims();
   PADDLE_ENFORCE_EQ(
-      (num_dims.size() == 1) && (num_dims[0] == 1),
+      phi::product(num_dims),
       true,
-      phi::errors::InvalidArgument("The shape of Input(Num) must be [1],"
-                                   "but received input shape is [%s].",
-                                   num_dims));
+      phi::errors::InvalidArgument("The size of Input(Num) must be 1,"
+                                   "but received input size is %s.",
+                                   phi::product(num_dims)));
   auto b_dims = base.dims();
-  PADDLE_ENFORCE_EQ(
-      (b_dims.size() == 1) && (b_dims[0] == 1),
-      true,
-      phi::errors::InvalidArgument("The shape of Input(Base) must be [1],"
-                                   "but received input shape is [%s].",
-                                   b_dims));
+  PADDLE_ENFORCE_EQ(phi::product(b_dims),
+                    true,
+                    phi::errors::InvalidArgument(
+                        "The size of Input(Base) must be 1,"
+                        "but received input size is phi::product(b_dims).",
+                        phi::product(b_dims)));
   out->set_dims(phi::make_ddim({-1}));
   out->set_dtype(dtype);
 }

paddle/phi/infermeta/unary.cc

@@ -2260,37 +2260,47 @@ void NanmedianInferMeta(const MetaTensor& x,
       for (int64_t i = 0; i < x_rank; i++) {
         out_dim.push_back(1);
       }
-    } else {
-      out_dim.push_back(1);
     }
   } else {
-    std::vector<int64_t> cleaned_axis;
+    std::vector<int64_t> formated_axis;
     for (auto& axis : axis_list) {
+      if (x_rank == 0) {
+        PADDLE_ENFORCE_EQ(axis == 0 || axis == -1,
+                          true,
+                          phi::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 dimesion = %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 dimesion = %d. But received axis = %d.",
+                              x_rank,
+                              axis));
+      }
       if (axis < 0) axis += x_rank;
-
-      PADDLE_ENFORCE_LT(
-          axis,
-          x_rank,
-          errors::InvalidArgument(
-              "Attr(axis) value should be in range [-R, R-1], R is "
-              "the rank of Input(X). But received axis: %d, R: %d. "
-              "Current Input(X)'s shape is=[%s].",
-              axis,
-              x_rank,
-              x_dim));
-
       PADDLE_ENFORCE_EQ(
-          std::find(cleaned_axis.begin(), cleaned_axis.end(), axis),
-          cleaned_axis.end(),
+          std::find(formated_axis.begin(), formated_axis.end(), axis),
+          formated_axis.end(),
           errors::InvalidArgument("Attr(axes) has duplicated elements: %d.",
                                   static_cast<int>(axis)));
 
-      cleaned_axis.push_back(axis);
+      formated_axis.push_back(axis);
     }
 
     for (int64_t i = 0; i < x_rank; i++) {
-      if (std::find(cleaned_axis.begin(), cleaned_axis.end(), i) ==
-          cleaned_axis.end()) {
+      if (std::find(formated_axis.begin(), formated_axis.end(), i) ==
+          formated_axis.end()) {
         out_dim.push_back(x_dim[i]);
       } else if (keep_dim) {
         out_dim.push_back(1);

paddle/phi/kernels/cpu/nanmedian_grad_kernel.cc

@@ -17,7 +17,7 @@
 #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/impl/nanmedian_grad_kernel_impl.h"
+#include "paddle/phi/kernels/funcs/nanmedian_utils.h"
 
 namespace phi {
 
@@ -26,67 +26,64 @@ void CalcMedianGradKernel(const Context& dev_ctx,
                           const DenseTensor& x,
                           const DenseTensor& median_index,
                           const DenseTensor& out_grad,
-                          const IntArray& axes UNUSED,
-                          DenseTensor* x_grad,
-                          T* x_grad_ptr) {
+                          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));
-  if (!x_grad_ptr) return;
 
-  const int64_t* m_ptr = median_index.data<int64_t>();
-  const T* out_grad_ptr = out_grad.data<T>();
+  const int64_t* m_data = median_index.data<int64_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[rank - 1];
-
   int64_t pre_dim = numel / stride;
+
   int64_t i = 0;
   int64_t offset = 0;
-  T div_factor = static_cast<T>(2.0);
   for (i = 0; i < pre_dim; i++) {
-    if (m_ptr[2 * i] >= 0) {
-      if (m_ptr[2 * i] == m_ptr[2 * i + 1]) {
-        x_grad_ptr[offset + m_ptr[2 * i]] = out_grad_ptr[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 {
-        x_grad_ptr[offset + m_ptr[2 * i]] = out_grad_ptr[i] / div_factor;
-        x_grad_ptr[offset + m_ptr[2 * i + 1]] = out_grad_ptr[i] / div_factor;
+        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, typename Context>
-void BaseMedianGradKernel(const Context& dev_ctx,
-                          const DenseTensor& x,
-                          const DenseTensor& median_index,
-                          const DenseTensor& out_grad,
-                          const IntArray& axes,
-                          DenseTensor* x_grad) {
-  auto rank = x.dims().size();
-  T* x_grad_ptr = dev_ctx.template Alloc<T>(x_grad);
-  if (axes.size() && (rank > 1)) {
-    DenseTensor tmp_x_grad(*x_grad);
-    CalcMedianGradKernel<T, Context>(
-        dev_ctx, x, median_index, out_grad, axes, &tmp_x_grad, x_grad_ptr);
-    PostprocessMedianGradKernel<T, Context>(dev_ctx, &tmp_x_grad, axes, x_grad);
-  } else {
-    CalcMedianGradKernel<T, Context>(
-        dev_ctx, x, median_index, out_grad, axes, x_grad, x_grad_ptr);
-  }
-}
-
 template <typename T, typename Context>
 void NanmedianGradKernel(const Context& dev_ctx,
-                         const DenseTensor& input,
+                         const DenseTensor& x,
                          const DenseTensor& median_index,
                          const DenseTensor& out_grad,
                          const IntArray& axes,
-                         bool keep_dim UNUSED,
+                         bool keepdim UNUSED,
                          DenseTensor* x_grad) {
-  BaseMedianGradKernel<T, Context>(
-      dev_ctx, input, median_index, out_grad, axes, 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);
+  } 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.template Alloc<T>(x_grad);
+    funcs::PostprocessMedianGradKernel<T, Context>(
+        dev_ctx, &tmp_x_grad, axes, x_grad);
+  }
 }
 
 }  // namespace phi

paddle/phi/kernels/cpu/nanmedian_kernel.cc

@@ -16,7 +16,7 @@
 
 #include "paddle/phi/backends/cpu/cpu_context.h"
 #include "paddle/phi/core/kernel_registry.h"
-#include "paddle/phi/kernels/impl/nanmedian_kernel_impl.h"
+#include "paddle/phi/kernels/funcs/nanmedian_utils.h"
 #include "paddle/phi/kernels/top_k_kernel.h"
 
 namespace phi {
@@ -31,7 +31,6 @@ void CalcMedianFunc(const Context& dev_ctx,
                     int64_t pre_dim,
                     T* o_ptr,
                     int64_t* m_ptr) {
-  bool should_ignore_nan = ignore_nan;
   DenseTensor sort_out;
   DenseTensor sort_indices;
   auto sort_dim = x.dims();
@@ -52,7 +51,7 @@ void CalcMedianFunc(const Context& dev_ctx,
   int64_t offset = 0;
   int64_t i = 0;
   bool is_ori_odd = stride & 1;
-  if (should_ignore_nan) {
+  if (ignore_nan) {
     for (i = 0; i < pre_dim; i++) {
       offset = i * sort_k;
       if (nan_counts[i] == stride) {
@@ -107,11 +106,11 @@ void CalcMedianFunc(const Context& dev_ctx,
 template <typename T, typename Context>
 void ProcessMedianKernel(const Context& dev_ctx,
                          const DenseTensor& x,
-                         T* o_ptr,
-                         int64_t* m_ptr,
-                         bool ignore_nan) {
-  bool should_ignore_nan = ignore_nan;
-  const T* x_ptr = x.data<T>();
+                         DenseTensor* out,
+                         DenseTensor* median_index) {
+  const T* x_data = x.data<T>();
+  T* out_data = dev_ctx.template Alloc<T>(out);
+  int64_t* m_data = dev_ctx.template Alloc<int64_t>(median_index);
 
   int64_t numel = x.numel();
   auto x_dim = x.dims();
@@ -122,7 +121,8 @@ void ProcessMedianKernel(const Context& dev_ctx,
 
   int64_t max_valid_num = 0;
   std::vector<int64_t> nan_counts;
-  if (should_ignore_nan) {
+  bool ignore_nan = true;
+  if (ignore_nan) {
     int64_t total_nan_num = 0;
     std::vector<T> col_vec;
     col_vec.reserve(stride);
@@ -133,7 +133,7 @@ void ProcessMedianKernel(const Context& dev_ctx,
     for (int64_t i = 0; i < pre_dim; i++) {
       col_vec.clear();
       col_vec.insert(
-          col_vec.begin(), x_ptr + i * stride, x_ptr + (i + 1) * stride);
+          col_vec.begin(), x_data + i * stride, x_data + (i + 1) * stride);
       nan_counts[i] =
           std::count_if(col_vec.begin(), col_vec.end(), [&](const T& val) {
             return std::isnan(static_cast<float>(val));
@@ -145,47 +145,25 @@ void ProcessMedianKernel(const Context& dev_ctx,
     // all elems are nan
     if (total_nan_num == numel) {
       for (i = 0; i < pre_dim; i++) {
-        o_ptr[i] = x_ptr[0];
-        m_ptr[2 * i] = -1;
-        m_ptr[2 * i + 1] = -1;
+        out_data[i] = std::numeric_limits<T>::quiet_NaN();
+        m_data[2 * i] = -1;
+        m_data[2 * i + 1] = -1;
       }
       return;
     }
-    should_ignore_nan = total_nan_num > 0;
+    ignore_nan = total_nan_num > 0;
   }
 
-  int64_t sort_k = should_ignore_nan ? max_valid_num : ((stride >> 1) + 1);
+  int64_t sort_k = ignore_nan ? max_valid_num : ((stride >> 1) + 1);
   CalcMedianFunc<T, Context>(dev_ctx,
                              x,
                              nan_counts,
-                             should_ignore_nan,
+                             ignore_nan,
                              sort_k,
                              stride,
                              pre_dim,
-                             o_ptr,
-                             m_ptr);
-}
-
-template <typename T, typename Context>
-void BaseMedianKernel(const Context& dev_ctx,
-                      const DenseTensor& input,
-                      const IntArray& axes,
-                      DenseTensor* out,
-                      DenseTensor* median_index,
-                      bool ignore_nan) {
-  DenseTensor x;
-  auto rank = input.dims().size();
-  if ((axes.size() == 0) || rank <= 1) {
-    x = input;
-    x.Resize({input.numel()});
-  } else {
-    PreprocessMedianKernel<T, Context>(dev_ctx, input, axes, &x);
-  }
-
-  T* o_ptr = dev_ctx.template Alloc<T>(out);
-  int64_t* m_ptr = dev_ctx.template Alloc<int64_t>(median_index);
-  ProcessMedianKernel<T, Context>(dev_ctx, x, o_ptr, m_ptr, ignore_nan);
-  out->Resize(out->dims());
+                             out_data,
+                             m_data);
 }
 
 template <typename T, typename Context>
@@ -195,7 +173,16 @@ void NanmedianKernel(const Context& dev_ctx,
                      bool keepdim UNUSED,
                      DenseTensor* out,
                      DenseTensor* median_index) {
-  BaseMedianKernel<T, Context>(dev_ctx, x, axes, out, median_index, true);
+  DenseTensor tmp_x;
+  auto rank = x.dims().size();
+  if ((axes.size() == 0) || rank <= 1) {
+    tmp_x = x;
+    tmp_x.Resize({x.numel()});
+  } else {
+    funcs::PreprocessMedianKernel<T, Context>(dev_ctx, x, axes, &tmp_x);
+  }
+
+  ProcessMedianKernel<T, Context>(dev_ctx, tmp_x, out, median_index);
 }
 
 }  // namespace phi

paddle/phi/kernels/impl/nanmedian_kernel_impl.h → paddle/phi/kernels/funcs/nanmedian_utils.h

@@ -15,9 +15,51 @@
 #pragma once
 
 #include "paddle/phi/kernels/funcs/math_function.h"
-#include "paddle/phi/kernels/nanmedian_kernel.h"
 
 namespace phi {
+namespace funcs {
+
+template <typename T, typename Context>
+void PostprocessMedianGradKernel(const Context& dev_ctx,
+                                 DenseTensor* input,
+                                 const IntArray& raw_axes,
+                                 DenseTensor* x) {
+  auto input_dim = input->dims();
+  auto rank = input_dim.size();
+
+  std::vector<int64_t> axes = raw_axes.GetData();
+  int64_t axes_size = static_cast<int>(axes.size());
+  for (int64_t i = 0; i < axes_size; i++) {
+    if (axes[i] < 0) {
+      axes[i] += rank;
+    }
+  }
+
+  std::vector<int> trans_back;
+  std::vector<int> reshape_back;
+  trans_back.resize(rank);
+
+  int offset = 0;
+  for (int64_t i = 0; i < rank; i++) {
+    if (std::find(axes.begin(), axes.end(), i) == axes.end()) {
+      reshape_back.push_back(input_dim[i]);
+      trans_back[i] = offset;
+      offset += 1;
+    }
+  }
+
+  for (int64_t i = 0; i < rank; i++) {
+    if (std::find(axes.begin(), axes.end(), i) != axes.end()) {
+      trans_back[i] = offset;
+      reshape_back.push_back(input_dim[i]);
+      offset += 1;
+    }
+  }
+
+  input->Resize(make_ddim(reshape_back));
+  funcs::TransCompute<Context, T>(
+      static_cast<int>(trans_back.size()), dev_ctx, *input, x, trans_back);
+}
 
 template <typename T, typename Context>
 void PreprocessMedianKernel(const Context& dev_ctx,
@@ -65,4 +107,5 @@ void PreprocessMedianKernel(const Context& dev_ctx,
   x->Resize(make_ddim(reshape));
 }
 
+}  // namespace funcs
 }  // namespace phi

paddle/phi/kernels/gpu/nanmedian_grad_kernel.cu

@@ -20,7 +20,7 @@
 #include "paddle/phi/core/kernel_registry.h"
 #include "paddle/phi/core/tensor_meta.h"
 #include "paddle/phi/kernels/funcs/math_function.h"
-#include "paddle/phi/kernels/impl/nanmedian_grad_kernel_impl.h"
+#include "paddle/phi/kernels/funcs/nanmedian_utils.h"
 
 namespace phi {
 
@@ -30,23 +30,26 @@ inline int GET_BLOCKS(const int N) {
 }
 
 template <typename T>
-__global__ void KernelNanmedianGrad(const T* x_ptr,
+__global__ void KernelNanmedianGrad(const T* x_data,
                                     const int64_t* medians_ptr,
                                     const T* out_grad_ptr,
-                                    T* x_grad_ptr,
+                                    T* dx_data,
                                     int64_t stride,
-                                    int64_t pre_dim,
-                                    T div_factor) {
+                                    int64_t pre_dim) {
   CUDA_KERNEL_LOOP(index, pre_dim) {
     int64_t offset = index * stride;
+    printf("index: %d\n", index);
+    printf("medians_ptr[2 * index]: %d\n", medians_ptr[2 * index]);
+    printf("medians_ptr[2 * index+1]: %d\n", medians_ptr[2 * index + 1]);
+
     if (medians_ptr[2 * index] >= 0) {
       if (medians_ptr[2 * index] == medians_ptr[2 * index + 1]) {
-        x_grad_ptr[offset + medians_ptr[2 * index]] = out_grad_ptr[index];
+        dx_data[offset + medians_ptr[2 * index]] = out_grad_ptr[index];
       } else {
-        x_grad_ptr[offset + medians_ptr[2 * index]] =
-            out_grad_ptr[index] / div_factor;
-        x_grad_ptr[offset + medians_ptr[2 * index + 1]] =
-            out_grad_ptr[index] / div_factor;
+        dx_data[offset + medians_ptr[2 * index]] =
+            out_grad_ptr[index] / static_cast<T>(2.0);
+        dx_data[offset + medians_ptr[2 * index + 1]] =
+            out_grad_ptr[index] / static_cast<T>(2.0);
       }
     }
   }
@@ -57,14 +60,16 @@ void CalcMedianGradKernel(const Context& dev_ctx,
                           const DenseTensor& x,
                           const DenseTensor& median_index,
                           const DenseTensor& out_grad,
-                          DenseTensor* x_grad,
-                          T* x_grad_ptr) {
+                          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));
 
   auto stream = dev_ctx.stream();
-  const T* x_ptr = x.data<T>();
-  const int64_t* m_ptr = median_index.data<int64_t>();
+  const T* x_data = x.data<T>();
+  const int64_t* m_data = median_index.data<int64_t>();
   const T* out_grad_ptr = out_grad.data<T>();
 
   int64_t numel = x.numel();
@@ -73,42 +78,38 @@ void CalcMedianGradKernel(const Context& dev_ctx,
   int64_t stride = x_dim[x_rank - 1];
   int64_t pre_dim = numel / stride;
 
-  T div_factor = static_cast<T>(2.0);
   KernelNanmedianGrad<T>
       <<<GET_BLOCKS(pre_dim), PADDLE_CUDA_NUM_THREADS, 0, stream>>>(
-          x_ptr, m_ptr, out_grad_ptr, x_grad_ptr, stride, pre_dim, div_factor);
-}
-
-template <typename T, typename Context>
-void BaseMedianGradKernel(const Context& dev_ctx,
-                          const DenseTensor& x,
-                          const DenseTensor& median_index,
-                          const DenseTensor& out_grad,
-                          const IntArray& axes,
-                          DenseTensor* x_grad) {
-  auto rank = x.dims().size();
-  T* x_grad_ptr = dev_ctx.template Alloc<T>(x_grad);
-  if (axes.size() && (rank > 1)) {
-    DenseTensor tmp_x_grad(*x_grad);
-    CalcMedianGradKernel<T, Context>(
-        dev_ctx, x, median_index, out_grad, &tmp_x_grad, x_grad_ptr);
-    PostprocessMedianGradKernel<T, Context>(dev_ctx, &tmp_x_grad, axes, x_grad);
-  } else {
-    CalcMedianGradKernel<T, Context>(
-        dev_ctx, x, median_index, out_grad, x_grad, x_grad_ptr);
-  }
+          x_data, m_data, out_grad_ptr, dx_data, stride, pre_dim);
 }
 
 template <typename T, typename Context>
 void NanmedianGradKernel(const Context& dev_ctx,
-                         const DenseTensor& input,
+                         const DenseTensor& x,
                          const DenseTensor& median_index,
                          const DenseTensor& out_grad,
                          const IntArray& axes,
-                         bool keep_dim,
+                         bool keepdim UNUSED,
                          DenseTensor* x_grad) {
-  BaseMedianGradKernel<T, Context>(
-      dev_ctx, input, median_index, out_grad, axes, 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);
+  } 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.template Alloc<T>(x_grad);
+    funcs::PostprocessMedianGradKernel<T, Context>(
+        dev_ctx, &tmp_x_grad, axes, x_grad);
+  }
 }
 
 }  // namespace phi

paddle/phi/kernels/gpu/nanmedian_kernel.cu

@@ -20,7 +20,7 @@
 #include "paddle/phi/common/memory_utils.h"
 #include "paddle/phi/core/kernel_registry.h"
 #include "paddle/phi/kernels/full_kernel.h"
-#include "paddle/phi/kernels/impl/nanmedian_kernel_impl.h"
+#include "paddle/phi/kernels/funcs/nanmedian_utils.h"
 #include "paddle/phi/kernels/top_k_kernel.h"
 
 namespace phi {
@@ -138,14 +138,13 @@ __global__ void CalcNanmedianKernel(const T* sort_out_ptr,
 template <typename T, typename Context>
 void ProcessMedianKernel(const Context& dev_ctx,
                          const DenseTensor& x,
-                         bool ignore_nan,
                          DenseTensor* out,
-                         int64_t* m_ptr) {
-  bool should_ignore_nan = ignore_nan;
+                         DenseTensor* median_index) {
   auto stream = dev_ctx.stream();
+  const T* x_data = x.data<T>();
+  T* out_data = dev_ctx.template Alloc<T>(out);
+  int64_t* m_data = dev_ctx.template Alloc<int64_t>(median_index);
 
-  const T* x_ptr = x.data<T>();
-  T* o_ptr = dev_ctx.template Alloc<T>(out);
   int64_t numel = x.numel();
   auto x_dim = x.dims();
   int64_t x_rank = x_dim.size();
@@ -156,7 +155,9 @@ void ProcessMedianKernel(const Context& dev_ctx,
   DenseTensor nan_counts, nan_stat;
   int64_t* nan_counts_ptr;
   int64_t max_valid_num = 0;
-  if (should_ignore_nan) {
+
+  bool ignore_nan = true;
+  if (ignore_nan) {
     nan_counts.Resize(phi::make_ddim({pre_dim}));
     dev_ctx.template Alloc<int64_t>(&nan_counts);
     nan_counts_ptr = nan_counts.data<int64_t>();
@@ -167,7 +168,7 @@ void ProcessMedianKernel(const Context& dev_ctx,
     KernelNanCounts<T><<<GET_BLOCKS(numel),
                          PADDLE_CUDA_NUM_THREADS,
                          pre_dim * sizeof(int64_t),
-                         stream>>>(x_ptr,
+                         stream>>>(x_data,
                                    numel,
                                    pre_dim,
                                    stride,
@@ -189,15 +190,19 @@ void ProcessMedianKernel(const Context& dev_ctx,
     // all elements are nan values
     T nan_val = std::numeric_limits<T>::quiet_NaN();
     if (nan_stat_cpu_ptr[0] == numel) {
-      FullLikeKernel<T, Context>(dev_ctx, x, nan_val, x.dtype(), out);
+      phi::funcs::SetConstant<Context, T> set_nan;
+      set_nan(dev_ctx, out, nan_val);
+
+      phi::funcs::SetConstant<Context, int64_t> set_negatvie;
+      set_negatvie(dev_ctx, median_index, static_cast<int64_t>(-1));
       return;
     }
 
-    should_ignore_nan = nan_stat_cpu_ptr[0] > 0;
+    ignore_nan = nan_stat_cpu_ptr[0] > 0;
     max_valid_num = nan_stat_cpu_ptr[1];
   }
 
-  int64_t sort_k = should_ignore_nan ? max_valid_num : ((stride >> 1) + 1);
+  int64_t sort_k = ignore_nan ? max_valid_num : ((stride >> 1) + 1);
   bool is_ori_odd = stride & 1;
 
   DenseTensor sort_out, sort_indices;
@@ -217,14 +222,14 @@ void ProcessMedianKernel(const Context& dev_ctx,
 
   T div_factor = static_cast<T>(2.0);
   T nan_val = std::numeric_limits<T>::quiet_NaN();
-  if (should_ignore_nan) {
+  if (ignore_nan) {
     CalcNanmedianKernel<T>
         <<<GET_BLOCKS(pre_dim), PADDLE_CUDA_NUM_THREADS, 0, stream>>>(
             sort_out_ptr,
             sort_indices_ptr,
             nan_counts_ptr,
-            m_ptr,
-            o_ptr,
+            m_data,
+            out_data,
             is_ori_odd,
             pre_dim,
             max_valid_num,
@@ -236,8 +241,8 @@ void ProcessMedianKernel(const Context& dev_ctx,
         <<<GET_BLOCKS(pre_dim), PADDLE_CUDA_NUM_THREADS, 0, stream>>>(
             sort_out_ptr,
             sort_indices_ptr,
-            m_ptr,
-            o_ptr,
+            m_data,
+            out_data,
             div_factor,
             is_ori_odd,
             pre_dim,
@@ -245,27 +250,6 @@ void ProcessMedianKernel(const Context& dev_ctx,
   }
 }
 
-template <typename T, typename Context>
-void BaseMedianKernel(const Context& dev_ctx,
-                      const DenseTensor& input,
-                      const IntArray& axes,
-                      bool ignore_nan,
-                      DenseTensor* out,
-                      DenseTensor* median_index) {
-  DenseTensor x;
-  auto rank = input.dims().size();
-  if ((axes.size() == 0) || rank <= 1) {
-    x = input;
-    x.Resize({input.numel()});
-  } else {
-    PreprocessMedianKernel<T, Context>(dev_ctx, input, axes, &x);
-  }
-
-  int64_t* m_ptr = dev_ctx.template Alloc<int64_t>(median_index);
-  ProcessMedianKernel<T, Context>(dev_ctx, x, ignore_nan, out, m_ptr);
-  out->Resize(out->dims());
-}
-
 template <typename T, typename Context>
 void NanmedianKernel(const Context& dev_ctx,
                      const DenseTensor& x,
@@ -273,7 +257,16 @@ void NanmedianKernel(const Context& dev_ctx,
                      bool keepdim,
                      DenseTensor* out,
                      DenseTensor* median_index) {
-  BaseMedianKernel<T, Context>(dev_ctx, x, axes, true, out, median_index);
+  DenseTensor tmp_x;
+  auto rank = x.dims().size();
+  if ((axes.size() == 0) || rank <= 1) {
+    tmp_x = x;
+    tmp_x.Resize({x.numel()});
+  } else {
+    funcs::PreprocessMedianKernel<T, Context>(dev_ctx, x, axes, &tmp_x);
+  }
+
+  ProcessMedianKernel<T, Context>(dev_ctx, tmp_x, out, median_index);
 }
 
 }  // namespace phi

paddle/phi/kernels/impl/nanmedian_grad_kernel_impl.h

-// Copyright (c) 2022 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.
-
-#pragma once
-
-#include "paddle/phi/kernels/funcs/math_function.h"
-#include "paddle/phi/kernels/nanmedian_grad_kernel.h"
-
-namespace phi {
-
-template <typename T, typename Context>
-void PostprocessMedianGradKernel(const Context& dev_ctx,
-                                 DenseTensor* input,
-                                 const IntArray& raw_axes,
-                                 DenseTensor* x) {
-  auto input_dim = input->dims();
-  auto rank = input_dim.size();
-
-  std::vector<int64_t> axes = raw_axes.GetData();
-  int64_t axes_size = static_cast<int>(axes.size());
-  for (int64_t i = 0; i < axes_size; i++) {
-    if (axes[i] < 0) {
-      axes[i] += rank;
-    }
-  }
-
-  std::vector<int> trans_back;
-  std::vector<int> reshape_back;
-  trans_back.reserve(rank);
-  trans_back.resize(rank);
-
-  int offset = 0;
-  for (int64_t i = 0; i < rank; i++) {
-    if (std::find(axes.begin(), axes.end(), i) == axes.end()) {
-      reshape_back.push_back(input_dim[i]);
-      trans_back[i] = offset;
-      offset += 1;
-    }
-  }
-
-  for (int64_t i = 0; i < rank; i++) {
-    if (std::find(axes.begin(), axes.end(), i) != axes.end()) {
-      trans_back[i] = offset;
-      reshape_back.push_back(input_dim[i]);
-      offset += 1;
-    }
-  }
-
-  input->Resize(make_ddim(reshape_back));
-  funcs::TransCompute<Context, T>(
-      static_cast<int>(trans_back.size()), dev_ctx, *input, x, trans_back);
-}
-
-}  // namespace phi

python/paddle/nn/functional/distance.py

@@ -64,8 +64,8 @@ def pairwise_distance(x, y, p=2.0, epsilon=1e-6, keepdim=False, name=None):
             y = paddle.to_tensor([[5., 6.], [7., 8.]], dtype=paddle.float64)
             distance = paddle.nn.functional.pairwise_distance(x, y)
             print(distance)
-    #       Tensor(shape=[2], dtype=float64, place=Place(gpu:0), stop_gradient=True,
-    #              [4.99999860, 4.99999860])
+            # Tensor(shape=[2], dtype=float64, place=Place(gpu:0), stop_gradient=True,
+            #        [4.99999860, 4.99999860])
 
     """
     if in_dynamic_mode():

python/paddle/tensor/creation.py

@@ -394,15 +394,15 @@ def logspace(start, stop, num, base=10.0, dtype=None, name=None):
 
     Args:
         start(int|float|Tensor): The input :attr:`start` is exponent of first entry in \
-            the sequence. It is a scalar, or a Tensor of shape [1] with input data \
+            the sequence. It is a scalar, or a 0-D Tensor of shape [] with input data \
             type int32, int64, float32 or float64.
         stop(int|float|Tensor): The input :attr:`stop` is exponent of last entry in the \
-            sequence. It is a scalar, or a Tensor of shape [1] with input data \
+            sequence. It is a scalar, or a 0-D Tensor of shape [] with input data \
             type int32, int64, float32 or float64.
         num(int|Tensor): The input :attr:`num` is given number of items in the sequence. \
-            It is an int scalar, or a Tensor of shape [1] with data type int32.
+            It is an int scalar, or a 0-D Tensor of shape [] with data type int32.
         base(int|float|Tensor): The input :attr:`base` is base of the logarithm function. \
-            It is a scalar, or a Tensor of shape [1] with input data type int32, int64, \
+            It is a scalar, or a 0-D Tensor of shape [] with input data type int32, int64, \
             float32 or float64.
         dtype(np.dtype|str, optional): The data type of output tensor, it could be \
             int32, int64, float32 or float64. Default: if None, the data type is float32. \

python/paddle/tensor/math.py

@@ -1615,7 +1615,7 @@ def count_nonzero(x, axis=None, keepdim=False, name=None):
             # x is a 2-D Tensor:
             x = paddle.to_tensor([[0., 1.1, 1.2], [0., 0., 1.3], [0., 0., 0.]])
             out1 = paddle.count_nonzero(x)
-            # [3]
+            # 3
             out2 = paddle.count_nonzero(x, axis=0)
             # [0, 1, 2]
             out3 = paddle.count_nonzero(x, axis=0, keepdim=True)
@@ -1636,17 +1636,8 @@ def count_nonzero(x, axis=None, keepdim=False, name=None):
             # [1, 3, 5]
     """
 
-    if axis is not None:
-        if isinstance(axis, int):
-            axis = [axis]
-        dims = len(x.shape)
-        for i in range(len(axis)):
-            if not isinstance(axis[i], int) or not (
-                axis[i] < dims and axis[i] >= -dims
-            ):
-                raise ValueError(
-                    "Axis should be None, int, or a list, element should in range [-rank(x), rank(x))."
-                )
+    if isinstance(axis, int):
+        axis = [axis]
 
     bool_tensor = paddle.cast(x, 'bool')
     int_tensor = paddle.cast(bool_tensor, 'int64')

python/paddle/tensor/stat.py

@@ -255,7 +255,7 @@ def numel(x, name=None):
         return out
 
 
-def nanmedian(x, axis=None, keepdim=True, name=None):
+def nanmedian(x, axis=None, keepdim=False, name=None):
     r"""
     Compute the median along the specified axis, while ignoring NaNs.
 
@@ -273,7 +273,7 @@ def nanmedian(x, axis=None, keepdim=True, name=None):
             in the output Tensor. If ``keepdim`` is True, the dimensions of
             the output Tensor is the same as ``x`` except in the reduced
             dimensions(it is of size 1 in this case). Otherwise, the shape of
-            the output Tensor is squeezed in ``axis`` . Default is True.
+            the output Tensor is squeezed in ``axis`` . Default is False.
         name (str, optional): Name for the operation (optional, default is None).
             For more information, please refer to :ref:`api_guide_Name`.
 
@@ -287,16 +287,16 @@ def nanmedian(x, axis=None, keepdim=True, name=None):
             x = paddle.to_tensor([[float('nan'), 2. , 3. ], [0. , 1. , 2. ]])
 
             y1 = x.nanmedian()
-            # y1 is [[2.]]
+            # y1 is 2.
 
             y2 = x.nanmedian(0)
-            # y2 is [[0.,  1.5, 2.5]]
+            # y2 is [0., 1.5, 2.5]
 
-            y3 = x.nanmedian(0, keepdim=False)
-            # y3 is [0.,  1.5, 2.5]
+            y3 = x.nanmedian(0, keepdim=True)
+            # y3 is [[0.,  1.5, 2.5]]
 
             y4 = x.nanmedian((0, 1))
-            # y4 is [[2.]]
+            # y4 is 2.
     """
     if not isinstance(x, Variable):
         raise TypeError("In median, the input x should be a Tensor.")
@@ -304,7 +304,6 @@ def nanmedian(x, axis=None, keepdim=True, name=None):
     if isinstance(axis, (list, tuple)) and len(axis) == 0:
         raise ValueError("Axis list should not be empty.")
 
-    dims = len(x.shape)
     if axis is None:
         axis = []
     elif isinstance(axis, tuple):
@@ -312,24 +311,6 @@ def nanmedian(x, axis=None, keepdim=True, name=None):
     elif isinstance(axis, int):
         axis = [axis]
 
-    if not isinstance(axis, list):
-        raise ValueError(
-            "Axis should be None, int, or a list, element should in range [-rank(x), rank(x))."
-        )
-
-    for i in range(len(axis)):
-        if not isinstance(axis[i], int) or not (
-            axis[i] < dims and axis[i] >= -dims
-        ):
-            raise ValueError(
-                "Axis should be None, int, or a list, element should in range [-rank(x), rank(x))."
-            )
-        if axis[i] < 0:
-            axis[i] += dims
-
-    if len(axis) != len(set(axis)):
-        raise ValueError("Axis has duplicated elements.")
-
     if in_dynamic_mode():
         return _C_ops.nanmedian(x, axis, keepdim)
     else:

test/legacy_test/test_nanmedian.py

@@ -125,6 +125,7 @@ class TestNanmedian(unittest.TestCase):
                 pd_res = paddle.nanmedian(
                     paddle.to_tensor(data), keepdim=keep_dim
                 )
+                assert np_res.shape == pd_res.numpy().shape
                 np.testing.assert_allclose(
                     np_res, pd_res.numpy(), rtol=1e-05, equal_nan=True
                 )
@@ -187,6 +188,23 @@ class TestNanmedian(unittest.TestCase):
         x_np[0, :] = np.nan
         x_np[1, :3] = np.nan
         x_np[2, 3:] = np.nan
+
+        x_tensor = paddle.to_tensor(x_np, stop_gradient=False)
+        y = paddle.nanmedian(x_tensor, keepdim=True)
+        dx = paddle.grad(y, x_tensor)[0].numpy()
+
+        np_grad = np.zeros(shape)
+        np_grad[1, 3] = 0.5
+        np_grad[3, 2] = 0.5
+        np.testing.assert_allclose(np_grad, dx, rtol=1e-05, equal_nan=True)
+
+    def test_check_grad_axis(self):
+        paddle.disable_static(place=self.place)
+        shape = (4, 5)
+        x_np = np.random.uniform(-1, 1, shape).astype(np.float64)
+        x_np[0, :] = np.nan
+        x_np[1, :3] = np.nan
+        x_np[2, 3:] = np.nan
         x_np_sorted = np.sort(x_np)
         nan_counts = np.count_nonzero(np.isnan(x_np).astype(np.int32), axis=1)
         np_grad = np.zeros(shape)
@@ -205,10 +223,25 @@ class TestNanmedian(unittest.TestCase):
                     np_grad[i, j] = 1 if is_odd else 0.5
 
         x_tensor = paddle.to_tensor(x_np, stop_gradient=False)
-        y = paddle.nanmedian(x_tensor, axis=1, keepdim=True)
+        y = paddle.nanmedian(x_tensor, axis=1)
         dx = paddle.grad(y, x_tensor)[0].numpy()
         np.testing.assert_allclose(np_grad, dx, rtol=1e-05, equal_nan=True)
 
+    def test_check_grad_0d(self):
+        paddle.disable_static(place=self.place)
+        x = paddle.rand([])
+        x.stop_gradient = False
+        y = paddle.nanmedian(x)
+        y.backward()
+        self.assertEqual(x.grad.shape, [])
+        np.testing.assert_allclose(x.grad, np.array(1.0))
+
+        x = paddle.to_tensor(float('nan'), stop_gradient=False)
+        y = paddle.nanmedian(x)
+        y.backward()
+        self.assertEqual(x.grad.shape, [])
+        np.testing.assert_allclose(x.grad, np.array(0.0))
+
 
 if __name__ == "__main__":
     unittest.main()