add erfinv FP16 test and BF16 test

paddle/phi/kernels/gpu/erfinv_grad_kernel.cu

@@ -19,8 +19,16 @@
 #include "paddle/phi/kernels/erfinv_grad_kernel.h"
 
 #include "paddle/phi/backends/gpu/gpu_context.h"
+#include "paddle/phi/common/bfloat16.h"
+#include "paddle/phi/common/float16.h"
 #include "paddle/phi/core/kernel_registry.h"
 #include "paddle/phi/kernels/impl/erfinv_grad_kernel_impl.h"
 
-PD_REGISTER_KERNEL(
-    erfinv_grad, GPU, ALL_LAYOUT, phi::ErfinvGradKernel, float, double) {}
+PD_REGISTER_KERNEL(erfinv_grad,
+                   GPU,
+                   ALL_LAYOUT,
+                   phi::ErfinvGradKernel,
+                   float,
+                   double,
+                   phi::dtype::float16,
+                   phi::dtype::bfloat16) {}

paddle/phi/kernels/gpu/erfinv_kernel.cu

@@ -14,6 +14,8 @@
 
 #include "paddle/phi/kernels/erfinv_kernel.h"
 #include "paddle/phi/backends/gpu/gpu_context.h"
+#include "paddle/phi/common/bfloat16.h"
+#include "paddle/phi/common/float16.h"
 #include "paddle/phi/core/kernel_registry.h"
 #include "paddle/phi/kernels/funcs/elementwise_base.h"
 
@@ -24,6 +26,22 @@ struct ErfinvFunctor {
   HOSTDEVICE inline T operator()(const T x) const { return erfinv(x); }
 };
 
+template <>
+struct ErfinvFunctor<float16> {
+  HOSTDEVICE inline float16 operator()(const float16 x) const {
+    auto newx = static_cast<float>(x);
+    return static_cast<float16>(erfinv(newx));
+  }
+};
+
+template <>
+struct ErfinvFunctor<bfloat16> {
+  HOSTDEVICE inline bfloat16 operator()(const bfloat16 x) const {
+    auto newx = static_cast<float>(x);
+    return static_cast<bfloat16>(erfinv(newx));
+  }
+};
+
 template <typename T, typename Context>
 void ErfinvKernel(const Context& ctx, const DenseTensor& x, DenseTensor* out) {
   ctx.template Alloc<T>(out);
@@ -34,4 +52,11 @@ void ErfinvKernel(const Context& ctx, const DenseTensor& x, DenseTensor* out) {
 
 }  // namespace phi
 
-PD_REGISTER_KERNEL(erfinv, GPU, ALL_LAYOUT, phi::ErfinvKernel, float, double) {}
+PD_REGISTER_KERNEL(erfinv,
+                   GPU,
+                   ALL_LAYOUT,
+                   phi::ErfinvKernel,
+                   float,
+                   double,
+                   phi::dtype::float16,
+                   phi::dtype::bfloat16) {}

paddle/phi/kernels/impl/erfinv_grad_kernel_impl.h

@@ -15,7 +15,10 @@
 #pragma once
 
 #include "paddle/phi/core/dense_tensor.h"
+#include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/kernels/erfinv_grad_kernel.h"
 #include "paddle/phi/kernels/funcs/eigen/common.h"
+#include "paddle/phi/kernels/funcs/eigen/eigen_function.h"
 
 namespace phi {
 
@@ -29,7 +32,7 @@ void ErfinvGradKernel(const Context& ctx,
   auto eigen_dout = EigenVector<T>::Flatten(out_grad);
   auto eigen_dx = EigenVector<T>::Flatten(*x_grad);
   auto& place = *ctx.eigen_device();
-  constexpr T half_sqrt_pi = static_cast<T>(1 / M_2_SQRTPI);
+  const T half_sqrt_pi = static_cast<T>(1 / M_2_SQRTPI);
   eigen_dx.device(place) = half_sqrt_pi * eigen_dout * eigen_out.square().exp();
 }
 

python/paddle/fluid/tests/unittests/test_erfinv_op.py

@@ -15,7 +15,7 @@
 import unittest
 
 import numpy as np
-from eager_op_test import OpTest
+from eager_op_test import OpTest, convert_float_to_uint16
 from scipy.special import erfinv
 
 import paddle
@@ -110,5 +110,42 @@ def run(place):
             run(place)
 
 
+class TestErfinvFP16OP(TestErfinv):
+    def init_dtype(self):
+        self.dtype = np.float16
+
+
+@unittest.skipIf(
+    not paddle.fluid.core.is_compiled_with_cuda()
+    or not paddle.fluid.core.is_bfloat16_supported(
+        paddle.fluid.core.CUDAPlace(0)
+    ),
+    "core is not complied with CUDA and not support the bfloat16",
+)
+class TestErfinvBF16OP(OpTest):
+    def setUp(self):
+        self.op_type = "erfinv"
+        self.public_python_api = paddle.erfinv
+        self.python_api = paddle.erfinv
+        self.dtype = np.uint16
+        self.shape = [11, 17]
+        x = np.random.uniform(-1, 1, size=self.shape).astype(np.float32)
+        res_ref = erfinv(x).astype(np.float32)
+        self.inputs = {'X': convert_float_to_uint16(x)}
+        self.outputs = {'Out': convert_float_to_uint16(res_ref)}
+
+    def test_check_output(self):
+        place = paddle.fluid.core.CUDAPlace(0)
+        self.check_output_with_place(place)
+
+    def test_check_grad(self):
+        place = paddle.fluid.core.CUDAPlace(0)
+        self.check_grad_with_place(
+            place,
+            ['X'],
+            'Out',
+        )
+
+
 if __name__ == "__main__":
     unittest.main()

python/paddle/tensor/math.py

@@ -4374,7 +4374,9 @@ def erfinv(x, name=None):
     if in_dygraph_mode():
         return _C_ops.erfinv(x)
     else:
-        check_variable_and_dtype(x, 'x', ['float32', 'float64'], 'erfinv')
+        check_variable_and_dtype(
+            x, 'x', ['float32', 'float64', 'float16', 'uint16'], 'erfinv'
+        )
         helper = LayerHelper('erfinv', **locals())
         out = helper.create_variable_for_type_inference(dtype=x.dtype)
         helper.append_op(type='erfinv', inputs={'X': x}, outputs={'Out': out})