[Accuracy diff No.21] Fix accuracy diff for heaviside API (#72894)

ooooo-create · web-flow · commit d14b9c44c456 · 2025-05-27T10:21:42.000+08:00
diff --git a/paddle/phi/kernels/cpu/elementwise_kernel.cc b/paddle/phi/kernels/cpu/elementwise_kernel.cc
@@ -74,8 +74,15 @@ void HeavisideKernel(const Context& dev_ctx,
                      DenseTensor* out) {
   // allocate memory for out
   dev_ctx.template Alloc<T>(out);
-  funcs::ElementwiseCompute<funcs::ElementwiseHeavisideFunctor<T>, T>(
-      dev_ctx, x, y, funcs::ElementwiseHeavisideFunctor<T>(), out);
+  auto x_dims = x.dims();
+  auto y_dims = y.dims();
+  if (x_dims.size() >= y_dims.size()) {
+    funcs::ElementwiseCompute<funcs::ElementwiseHeavisideFunctor<T>, T>(
+        dev_ctx, x, y, funcs::ElementwiseHeavisideFunctor<T>(), out);
+  } else {
+    funcs::ElementwiseCompute<funcs::ElementwiseInverseHeavisideFunctor<T>, T>(
+        dev_ctx, x, y, funcs::ElementwiseInverseHeavisideFunctor<T>(), out);
+  }
 }
 
 template <typename T, typename Context>
diff --git a/paddle/phi/kernels/funcs/elementwise_functor.h b/paddle/phi/kernels/funcs/elementwise_functor.h
@@ -801,6 +801,13 @@ struct ElementwiseHeavisideFunctor {
   }
 };
 
+template <typename T>
+struct ElementwiseInverseHeavisideFunctor {
+  inline HOSTDEVICE T operator()(const T a, const T b) const {
+    return b == static_cast<T>(0) ? a : static_cast<T>(b > static_cast<T>(0));
+  }
+};
+
 template <typename T, typename Enable = void>
 struct FloorDivideFunctor {
   inline HOSTDEVICE T operator()(const T a, const T b) const {
diff --git a/test/legacy_test/test_elementwise_heaviside_op.py b/test/legacy_test/test_elementwise_heaviside_op.py
@@ -205,6 +205,30 @@ def setUp(self):
         self.outputs = {'Out': np.heaviside(self.inputs['X'], self.inputs['Y'])}
 
 
+class TestElementwiseOp3(TestElementwiseOp):
+    def setUp(self):
+        self.op_type = "elementwise_heaviside"
+        x = np.random.uniform(-10, 10, [100]).astype("float64")
+        y = np.random.uniform(-10, 10, [3, 100]).astype("float64")
+        self.python_api = paddle.heaviside
+        self.prim_op_type = "comp"
+        self.public_python_api = paddle.heaviside
+        self.inputs = {'X': x, 'Y': y}
+        self.outputs = {'Out': np.heaviside(self.inputs['X'], self.inputs['Y'])}
+
+
+class TestElementwiseOp4(TestElementwiseOp):
+    def setUp(self):
+        self.op_type = "elementwise_heaviside"
+        x = np.random.uniform(0, 10, []).astype("float64")
+        y = np.random.uniform(-10, 0, [2, 3, 20]).astype("float64")
+        self.python_api = paddle.heaviside
+        self.prim_op_type = "comp"
+        self.public_python_api = paddle.heaviside
+        self.inputs = {'X': x, 'Y': y}
+        self.outputs = {'Out': np.heaviside(self.inputs['X'], self.inputs['Y'])}
+
+
 class TestHeavisideFP16Op(OpTest):
     def setUp(self):
         self.dtype = np.float16
