fix unique_kernel support axis=-1 (#49385)

ab786715 · sprouteer · GitHub · da357615 · ab786715 · ab786715
隐藏空白更改
内联并排

Showing with 2 addition and 0 deletion

paddle/phi/kernels/cpu/unique_kernel.cc paddle/phi/kernels/cpu/unique_kernel.cc +1 -0

paddle/phi/kernels/gpu/unique_kernel.cu paddle/phi/kernels/gpu/unique_kernel.cu +1 -0

未找到文件。
--- 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,

--- 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,