Add support for np.asarray_chkfinite [WIP] [need help]

docs/source/reference/numpysupported.rst

@@ -338,6 +338,7 @@ The following top-level functions are supported:
 * :func:`numpy.array` (only the 2 first arguments)
 * :func:`numpy.array_equal`
 * :func:`numpy.asarray` (only the 2 first arguments)
+* :func:`numpy.asarray_chkfinite` (only the first first arguments)
 * :func:`numpy.asfarray` 
 * :func:`numpy.asfortranarray` (only the first argument)
 * :func:`numpy.atleast_1d`

numba/np/arraymath.py

@@ -4159,6 +4159,28 @@ def np_select_arr_impl(condlist, choicelist, default=0):
 
     return np_select_arr_impl
 
+
+@overload(np.asarray_chkfinite)
+def np_asarray_chkfinite(a, dtype=None, order='C'):
+
+    msg = "The argument to np.asarray_chkfinite must be array-like"
+    if not isinstance(a, (types.Array, types.Sequence, types.Tuple)):
+        raise TypingError(msg)
+
+    if is_nonelike(dtype):
+        dt = a.dtype
+    else:
+        dt = dtype.dtype
+
+    def impl(a, dtype=None, order='C'):
+        a = np.asarray(a, dtype=dt)
+        for i in np.nditer(a):
+            if not np.isfinite(i):
+                raise ValueError("array must not contain infs or NaNs")
+        return a
+
+    return impl
+
 #----------------------------------------------------------------------------
 # Windowing functions
 #   - translated from the numpy implementations found in:

numba/tests/test_np_functions.py

@@ -321,6 +321,10 @@ def flip_ud(a):
     return np.flipud(a)
 
 
+def np_asarray_chkfinite(a, dtype=None, order='C'):
+    return np.asarray_chkfinite(a, dtype, order)
+
+
 class TestNPFunctions(MemoryLeakMixin, TestCase):
     """
     Tests for various Numpy functions.
@@ -3707,6 +3711,82 @@ def test_cross2d_exceptions(self):
             str(raises.exception)
         )
 
+    def test_asarray_chkfinite(self):
+        pyfunc = np_asarray_chkfinite
+        cfunc = jit(nopython=True)(pyfunc)
+        self.disable_leak_check()
+
+        pairs = [
+            #1D array with all args
+            (
+                np.array([1, 2, 3]),
+                np.float32,
+                'C'
+            ),
+            #1D array
+            (
+                np.array([1, 2, 3]),
+            ),
+            #1D array-like
+            (
+                [1, 2, 3, 4],
+            ),
+            # 2x2 (n-dims)
+            (
+                np.array([[1, 2], [3, 4]]),
+                np.float32,
+                'C'
+            ),
+            # 2x2 array-like (n-dims)
+            (
+                ((1, 2), (3, 4)),
+                np.int64
+            ),
+            # 2x2 (1-dim) with type promotion
+            (
+                np.array([1, 2], dtype=np.int64),
+            ),
+            # 3x2 (with higher order broadcasting)
+            (
+                np.arange(36).reshape(6, 2, 3),
+            )
+        ]
+
+        for pair in pairs:
+            expected = pyfunc(*pair)
+            got = cfunc(*pair)
+            self.assertPreciseEqual(expected, got)
+
+    def test_asarray_chkfinite_exceptions(self):
+        cfunc = jit(nopython=True)(np_asarray_chkfinite)
+        self.disable_leak_check()
+
+        #test for single value
+        with self.assertRaises(TypingError) as e:
+            cfunc(2)
+        msg = "The argument to np.asarray_chkfinite must be array-like"
+        self.assertIn(msg, str(e.exception))
+
+        #test for NaNs
+        with self.assertRaises(ValueError) as e:
+            cfunc(np.array([2, 4, np.nan, 5]))
+        self.assertIn("array must not contain infs or NaNs", str(e.exception))
+
+        #test for infs
+        with self.assertRaises(ValueError) as e:
+            cfunc(np.array([1, 2, np.inf, 4]))
+        self.assertIn("array must not contain infs or NaNs", str(e.exception))
+
+        #test for both inf and NaNs
+        with self.assertRaises(ValueError) as e:
+            cfunc(np.array([np.inf, np.nan]))
+        self.assertIn("array must not contain infs or NaNs", str(e.exception))
+
+        #test for NaNs
+        with self.assertRaises(ValueError) as e:
+            cfunc(np.array([np.nan, np.nan, np.nan]))
+        self.assertIn("array must not contain infs or NaNs", str(e.exception))
+
 
 class TestNPMachineParameters(TestCase):
     # tests np.finfo, np.iinfo, np.MachAr