fix unique_kernel support axis=-1

PaddlePaddle · Dec 22, 2022 · ad9bf65 · ad9bf65
1 parent 866f492
commit ad9bf65
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 0 deletions.
diff --git a/paddle/phi/kernels/cpu/unique_kernel.cc b/paddle/phi/kernels/cpu/unique_kernel.cc
@@ -96,6 +96,7 @@ void UniqueRawKernel(const Context& context,
                                                             return_counts));
   } else {
     int axis_value = axis[0];
+    axis_value = (axis_value == -1) ? (x.dims().size() - 1) : axis_value;
     phi::VisitDataTypeTiny(
         dtype,
         phi::funcs::UniqueDimFunctor<Context, T>(context,

diff --git a/paddle/phi/kernels/gpu/unique_kernel.cu b/paddle/phi/kernels/gpu/unique_kernel.cu
@@ -563,6 +563,7 @@ void UniqueRawKernel(const Context& context,
   } else {
     // 'axis' is required.
     int axis_value = axis[0];
+    axis_value = (axis_value == -1) ? (x.dims().size() - 1) : axis_value;
     phi::VisitDataTypeTiny(dtype,
                            UniqueDimsCUDAFunctor<Context, T>(context,
                                                              x,