Repository: cupy/cupy

Issue:
Supporting nan_to_num's "nan", "posinf", and "neginf" arguments
[cupy.nan_to_num](https://docs.cupy.dev/en/stable/reference/generated/cupy.nan_to_num.html) was implemented in #1391, but as a ufunc instead of a function that takes non-array arguments, as in NumPy. This means we can't pass the `nan`, `posinf`, or `neginf` arguments to it, as in [numpy.nan_to_num](https://docs.scipy.org/doc/numpy/reference/generated/numpy.nan_to_num.html).

My CuPy version is 8.5.0 (updated to the latest on conda-forge).


Discussion:

Resolving patch:
diff --git a/cupy/_math/misc.py b/cupy/_math/misc.py
--- a/cupy/_math/misc.py
+++ b/cupy/_math/misc.py
@@ -316,46 +316,75 @@ def clip(a, a_min=None, a_max=None, out=None):
 
 _nan_to_num_preamble = '''
 template <class T>
-__device__ T nan_to_num(T x, T large) {
+__device__ T nan_to_num(T x, T nan, T posinf, T neginf) {
     if (isnan(x))
-        return 0;
+        return nan;
     if (isinf(x))
-        return copysign(large, x);
+        return x > 0 ? posinf : neginf;
     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);
+__device__ complex<T> nan_to_num(complex<T> x, T nan, T posinf, T neginf) {
+    T re = nan_to_num(x.real(), nan, posinf, neginf);
+    T im = nan_to_num(x.imag(), nan, posinf, neginf);
     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))')),
+_nan_to_num = _core.create_ufunc(
+    'cupy_nan_to_num_',
+    ('????->?', 'bbbb->b', 'BBBB->B', 'hhhh->h', 'HHHH->H',
+     'iiii->i', 'IIII->I', 'llll->l', 'LLLL->L', 'qqqq->q', 'QQQQ->Q',
+     ('eeee->e',
+      'out0 = nan_to_num(in0, in1, in2, in3)'),
+     ('ffff->f',
+      'out0 = nan_to_num(in0, in1, in2, in3)'),
+     ('dddd->d',
+      'out0 = nan_to_num(in0, in1, in2, in3)'),
+     ('Ffff->F',
+      'out0 = nan_to_num(in0, in1, in2, in3)'),
+     ('Dddd->D',
+      'out0 = nan_to_num(in0, in1, in2, in3)')),
     'out0 = in0',
     preamble=_nan_to_num_preamble,
-    doc='''Elementwise nan_to_num function.
+    doc=''' Elementwise nan_to_num function.
 
-    .. seealso:: :data:`numpy.nan_to_num`
+    .. seealso:: :func:`numpy.nan_to_num`
 
     ''')
 
 
+def _check_nan_inf(x, dtype, neg=None):
+    if dtype.char in 'FD':
+        dtype = cupy.dtype(dtype.char.lower())
+    if dtype.char not in 'efd':
+        x = 0
+    elif x is None and neg is not None:
+        x = cupy.finfo(dtype).min if neg else cupy.finfo(dtype).max
+    elif cupy.isnan(x):
+        x = cupy.nan
+    elif cupy.isinf(x):
+        x = cupy.inf * (-1)**(x < 0)
+    return cupy.asanyarray(x, dtype)
+
+
+def nan_to_num(x, out=None, *, nan=0.0, posinf=None, neginf=None, **kwds):
+    """ Elementwise nan_to_num function.
+
+    .. seealso:: :func:`numpy.nan_to_num`
+
+    """
+
+    kwds.setdefault('out', out)
+    dtype = cupy.asanyarray(x).dtype
+    nan = _check_nan_inf(nan, dtype)
+    posinf = _check_nan_inf(posinf, dtype, False)
+    neginf = _check_nan_inf(neginf, dtype, True)
+    return _nan_to_num(x, nan, posinf, neginf, **kwds)
+
+
 # TODO(okuta): Implement real_if_close
 
 


Test patch: