Repository: cupy/cupy

Issue:
Implement nan_to_num
🙏 
https://docs.scipy.org/doc/numpy/reference/generated/numpy.nan_to_num.html


Discussion:

Resolving patch:
diff --git a/cupy/__init__.py b/cupy/__init__.py
--- a/cupy/__init__.py
+++ b/cupy/__init__.py
@@ -546,6 +546,7 @@ def isscalar(num):
 from cupy.core.fusion import fmin  # NOQA
 from cupy.core.fusion import maximum  # NOQA
 from cupy.core.fusion import minimum  # NOQA
+from cupy.core.fusion import nan_to_num  # NOQA
 from cupy.core.fusion import sign  # NOQA
 from cupy.core.fusion import sqrt  # NOQA
 from cupy.core.fusion import square  # NOQA
diff --git a/cupy/core/fusion.py b/cupy/core/fusion.py
--- a/cupy/core/fusion.py
+++ b/cupy/core/fusion.py
@@ -963,6 +963,7 @@ def __call__(self, *args, **kwargs):
 minimum = ufunc(math.misc.minimum)
 fmax = ufunc(math.misc.fmax)
 fmin = ufunc(math.misc.fmin)
+nan_to_num = ufunc(math.misc.nan_to_num)
 
 if hasattr(numpy, "divmod"):
     divmod = ufunc(core.divmod)
diff --git a/cupy/math/misc.py b/cupy/math/misc.py
--- a/cupy/math/misc.py
+++ b/cupy/math/misc.py
@@ -152,7 +152,46 @@ def clip(a, a_min=None, a_max=None, out=None):
     ''')
 
 
-# TODO(okuta): Implement nan_to_num
+_nan_to_num_preamble = '''
+template <class T>
+__device__ T nan_to_num(T x, T large) {
+    if (isnan(x))
+        return 0;
+    if (isinf(x))
+        return copysign(large, x);
+    return x;
+}
+
+template <class T>
+__device__ complex<T> nan_to_num(complex<T> x, T large) {
+    T re = nan_to_num(x.real(), large);
+    T im = nan_to_num(x.imag(), large);
+    return complex<T>(re, im);
+}
+'''
+
+
+nan_to_num = core.create_ufunc(
+    'cupy_nan_to_num',
+    ('?->?', 'b->b', 'B->B', 'h->h', 'H->H',
+     'i->i', 'I->I', 'l->l', 'L->L', 'q->q', 'Q->Q',
+     ('e->e',
+      'out0 = nan_to_num(in0, float16(32 * 0x7FF))'),
+     ('f->f',
+      'out0 = nan_to_num(in0, __int_as_float(0x7F800000 - 1))'),
+     ('d->d',
+      'out0 = nan_to_num(in0, __longlong_as_double(0x7FF0000000000000 - 1))'),
+     ('F->F',
+      'out0 = nan_to_num(in0, __int_as_float(0x7F800000 - 1))'),
+     ('D->D',
+      'out0 = nan_to_num(in0, __longlong_as_double(0x7FF0000000000000 - 1))')),
+    'out0 = in0',
+    preamble=_nan_to_num_preamble,
+    doc='''Elementwise nan_to_num function.
+
+    .. seealso:: :data:`numpy.nan_to_num`
+
+    ''')
 
 
 # TODO(okuta): Implement real_if_close


Test patch:
diff --git a/tests/cupy_tests/math_tests/test_misc.py b/tests/cupy_tests/math_tests/test_misc.py
--- a/tests/cupy_tests/math_tests/test_misc.py
+++ b/tests/cupy_tests/math_tests/test_misc.py
@@ -1,3 +1,4 @@
+from distutils.version import StrictVersion
 import sys
 import unittest
 
@@ -36,6 +37,46 @@ def check_unary_negative(self, name, xp, dtype, no_bool=False):
             a += (a * 1j).astype(dtype)
         return getattr(xp, name)(a)
 
+    @testing.for_dtypes(['e', 'f', 'd', 'F', 'D'])
+    @testing.numpy_cupy_allclose(atol=1e-5)
+    def check_unary_inf(self, name, xp, dtype):
+        inf = numpy.inf
+        if numpy.dtype(dtype).kind != 'c':
+            a = xp.array([0, -1, 1, -inf, inf], dtype=dtype)
+        else:
+            a = xp.array([complex(x, y)
+                          for x in [0, -1, 1, -inf, inf]
+                          for y in [0, -1, 1, -inf, inf]],
+                         dtype=dtype)
+        return getattr(xp, name)(a)
+
+    @testing.for_dtypes(['e', 'f', 'd', 'F', 'D'])
+    @testing.numpy_cupy_allclose(atol=1e-5)
+    def check_unary_nan(self, name, xp, dtype):
+        nan = numpy.nan
+        if numpy.dtype(dtype).kind != 'c':
+            a = xp.array([0, -1, 1, -nan, nan], dtype=dtype)
+        else:
+            a = xp.array([complex(x, y)
+                          for x in [0, -1, 1, -nan, nan]
+                          for y in [0, -1, 1, -nan, nan]],
+                         dtype=dtype)
+        return getattr(xp, name)(a)
+
+    @testing.for_dtypes(['e', 'f', 'd', 'F', 'D'])
+    @testing.numpy_cupy_allclose(atol=1e-5)
+    def check_unary_inf_nan(self, name, xp, dtype):
+        inf = numpy.inf
+        nan = numpy.nan
+        if numpy.dtype(dtype).kind != 'c':
+            a = xp.array([0, -1, 1, -inf, inf, -nan, nan], dtype=dtype)
+        else:
+            a = xp.array([complex(x, y)
+                          for x in [0, -1, 1, -inf, inf, -nan, nan]
+                          for y in [0, -1, 1, -inf, inf, -nan, nan]],
+                         dtype=dtype)
+        return getattr(xp, name)(a)
+
     @testing.for_dtypes(['e', 'f', 'd', 'F', 'D'])
     @testing.numpy_cupy_array_equal()
     def check_binary_nan(self, name, xp, dtype):
@@ -147,3 +188,20 @@ def test_fmin(self):
 
     def test_fmin_nan(self):
         self.check_binary_nan('fmin')
+
+    def test_nan_to_num(self):
+        no_bool = StrictVersion(numpy.version.version) < StrictVersion('1.10')
+        self.check_unary('nan_to_num', no_bool=no_bool)
+
+    def test_nan_to_num_negative(self):
+        no_bool = StrictVersion(numpy.version.version) < StrictVersion('1.10')
+        self.check_unary_negative('nan_to_num', no_bool=no_bool)
+
+    def test_nan_to_num_inf(self):
+        self.check_unary_inf('nan_to_num')
+
+    def test_nan_to_num_nan(self):
+        self.check_unary_nan('nan_to_num')
+
+    def test_nan_to_num_inf_nan(self):
+        self.check_unary_inf_nan('nan_to_num')