Fix erfinv and swapaxes

oneflow/core/functional/impl/math_functor.cpp

@@ -648,9 +648,6 @@ class SwapaxesFunctor {
         << "], but got " << dim_1 << ")";
     return Transpose2dim(x, dim0, dim1);
   }
-
- private:
-  std::shared_ptr<OpExpr> op_;
 };
 
 class ArangeFunctor {

oneflow/user/kernels/erfinv_kernel.cpp

@@ -14,7 +14,7 @@ See the License for the specific language governing permissions and
 limitations under the License.
 */
 #include "oneflow/core/framework/framework.h"
-#include <math.h>
+#define PI 3.14159265358979323846L
 namespace oneflow {
 
 template<typename T>
@@ -39,8 +39,14 @@ class CpuErfinvKernel final : public user_op::OpKernel {
       T z, num, dem;
       T x = x_ptr[i];  // Promise the correctness of inplace version.
       T x_abs = std::abs(x);
-      if (x_abs > 1.0) y_ptr[i] = std::numeric_limits<T>::quiet_NaN();
-      if (x_abs == 1.0) y_ptr[i] = std::copysign(std::numeric_limits<T>::infinity(), x);
+      if (x_abs > 1.0) {
+        y_ptr[i] = std::numeric_limits<T>::quiet_NaN();
+        continue;
+      }
+      if (x_abs == 1.0) {
+        y_ptr[i] = std::copysign(std::numeric_limits<T>::infinity(), x);
+        continue;
+      }
       if (x_abs <= static_cast<T>(central_range)) {
         z = x * x;
         num = (((a[3] * z + a[2]) * z + a[1]) * z + a[0]);
@@ -54,11 +60,11 @@ class CpuErfinvKernel final : public user_op::OpKernel {
       }
       y_ptr[i] = y_ptr[i]
                  - (std::erf(y_ptr[i]) - x)
-                       / ((static_cast<T>(2.0) / static_cast<T>(std::sqrt(M_PI)))
+                       / ((static_cast<T>(2.0) / static_cast<T>(std::sqrt(static_cast<double>(PI))))
                           * std::exp(-y_ptr[i] * y_ptr[i]));
       y_ptr[i] = y_ptr[i]
                  - (std::erf(y_ptr[i]) - x)
-                       / ((static_cast<T>(2.0) / static_cast<T>(std::sqrt(M_PI)))
+                       / ((static_cast<T>(2.0) / static_cast<T>(std::sqrt(static_cast<double>(PI))))
                           * std::exp(-y_ptr[i] * y_ptr[i]));
     }
   }

python/oneflow/test/modules/test_erfinv.py

@@ -27,8 +27,34 @@
 from oneflow.test_utils.automated_test_util import *
 
 
+def _test_flow_erfinv_with_inf_data(test_case, device):
+    x = flow.tensor(np.ones((5, 5)), dtype=flow.float32, device=flow.device(device))
+    of_out = flow.erfinv(x)
+    np_out = np.full((5, 5), np.inf)
+    test_case.assertTrue(np.array_equal(of_out.numpy(), np_out))
+
+
+def _test_flow_erfinv_with_nan_data(test_case, device):
+    x = flow.tensor(
+        np.arange(2, 22).reshape(4, 5), dtype=flow.float32, device=flow.device(device)
+    )
+    of_out = flow.erfinv(x)
+    np_out = np.full((4, 5), np.nan)
+    test_case.assertTrue(np.array_equal(of_out.numpy(), np_out, equal_nan=True))
+
+
 @flow.unittest.skip_unless_1n1d()
 class TestErfinvModule(flow.unittest.TestCase):
+    def test_flow_erfinv(test_case):
+        arg_dict = OrderedDict()
+        arg_dict["test_fun"] = [
+            _test_flow_erfinv_with_inf_data,
+            _test_flow_erfinv_with_nan_data,
+        ]
+        arg_dict["device"] = ["cpu", "cuda"]
+        for arg in GenArgList(arg_dict):
+            arg[0](test_case, *arg[1:])
+
     @autotest(check_graph=True, auto_backward=False)
     def test_flow_erfinv_with_random_data(test_case):
         device = random_device()