Repository: scikit-learn/scikit-learn
Issue category: Performance Issue

Optimizing assert_all_finite check
Have done some work to come up with a Cython implementation of [`_assert_all_finite`]( https://github.com/scikit-learn/scikit-learn/blob/6d15840432c08b2c628babe62d7bef6d6a01fcf6/sklearn/utils/validation.py#L40 ) as the current check is rather slow. Have come up with [this implementation]( https://gist.github.com/jakirkham/8bafcec842e92641885c91f79ef7f44f ). This simply uses the C `isfinite` (or `isinf` for `allow_nan=True`) function in a tight loop; returning `False` whenever `isfinite` returns `False` and returning `True` otherwise. As such it has a very small, constant memory overhead.

For smallish data (0-8000 element `float64` arrays), this is comparable to the naive NumPy implementation in time and faster than the current scikit-learn implementation. Should add that the Cython implementation uses considerably less memory than either implementation. For large sized arrays, scikit-learn's performance improves to ~2x faster than the NumPy or Cython implementation. Expect this is mostly due to how the caches are utilized. So this can be improved further.

Would be interested in getting some feedback about whether such a contribution would be interesting to scikit-learn. If so, can investigate further improvements to this implementation.


Functions to edit:
sklearn/utils/setup.py:configuration
sklearn/utils/validation.py:_assert_all_finite

Patch:
diff --git a/doc/whats_new/v1.2.rst b/doc/whats_new/v1.2.rst
index 832a4e4389b19..22aecf64a85b8 100644
--- a/doc/whats_new/v1.2.rst
+++ b/doc/whats_new/v1.2.rst
@@ -44,6 +44,11 @@ Changes impacting all modules
   NumPy's SIMD optimized primitives.
   :pr:`23446` by :user:`Meekail Zain <micky774>`
 
+- |Enhancement| Finiteness checks (detection of NaN and infinite values) in all
+  estimators are now faster by utilizing a more efficient stop-on-first
+  second-pass algorithm.
+  :pr:`23197` by :user:`Meekail Zain <micky774>`
+
 Changelog
 ---------
 
diff --git a/sklearn/utils/_isfinite.pyx b/sklearn/utils/_isfinite.pyx
new file mode 100644
index 0000000000000..bc6af6b8c24a1
--- /dev/null
+++ b/sklearn/utils/_isfinite.pyx
@@ -0,0 +1,40 @@
+# Author: jakirkham, Meekail Zain, Thomas Fan
+
+from libc.math cimport isnan, isinf
+cimport cython
+from cython cimport floating
+
+cpdef enum FiniteStatus:
+    all_finite = 0
+    has_nan = 1
+    has_infinite = 2
+
+def cy_isfinite(floating[::1] a, bint allow_nan=False):
+    cdef int result
+    with nogil:
+        result = _isfinite(a, allow_nan)
+    return result
+
+cdef inline int _isfinite(floating[::1] a, bint allow_nan) nogil:
+    cdef floating* a_ptr = &a[0]
+    cdef Py_ssize_t length = len(a)
+    if allow_nan:
+        return _isfinite_allow_nan(a_ptr, length)
+    else:
+        return _isfinite_disable_nan(a_ptr, length)
+
+cdef inline int _isfinite_allow_nan(floating* a_ptr, Py_ssize_t length) nogil:
+    cdef Py_ssize_t i
+    for i in range(length):
+        if isinf(a_ptr[i]) != 0:
+            return FiniteStatus.has_infinite
+    return FiniteStatus.all_finite
+
+cdef inline int _isfinite_disable_nan(floating* a_ptr, Py_ssize_t length) nogil:
+    cdef Py_ssize_t i
+    for i in range(length):
+        if isnan(a_ptr[i]) != 0:
+            return FiniteStatus.has_nan
+        if isinf(a_ptr[i]) != 0:
+            return FiniteStatus.has_infinite
+    return FiniteStatus.all_finite
diff --git a/sklearn/utils/setup.py b/sklearn/utils/setup.py
index eaf3a447389ea..915a8efeb2e01 100644
--- a/sklearn/utils/setup.py
+++ b/sklearn/utils/setup.py
@@ -116,6 +116,11 @@ def configuration(parent_package="", top_path=None):
         libraries=libraries,
         language="c++",
     )
+    config.add_extension(
+        "_isfinite",
+        sources=["_isfinite.pyx"],
+        libraries=libraries,
+    )
 
     config.add_subpackage("tests")
 
diff --git a/sklearn/utils/validation.py b/sklearn/utils/validation.py
index bb2c55d4beed5..6bfe99fdaffdc 100644
--- a/sklearn/utils/validation.py
+++ b/sklearn/utils/validation.py
@@ -29,6 +29,7 @@
 from ..exceptions import PositiveSpectrumWarning
 from ..exceptions import NotFittedError
 from ..exceptions import DataConversionWarning
+from ._isfinite import cy_isfinite, FiniteStatus
 
 FLOAT_DTYPES = (np.float64, np.float32, np.float16)
 
@@ -99,25 +100,33 @@ def _assert_all_finite(
         return
     X = np.asanyarray(X)
     # First try an O(n) time, O(1) space solution for the common case that
-    # everything is finite; fall back to O(n) space np.isfinite to prevent
-    # false positives from overflow in sum method.
+    # everything is finite; fall back to O(n) space `np.isinf/isnan` or custom
+    # Cython implementation to prevent false positives and provide a detailed
+    # error message.
     is_float = X.dtype.kind in "fc"
-    with np.errstate(over="ignore"):
-        first_pass_isfinite = is_float and np.isfinite(np.sum(X))
-    if first_pass_isfinite:
-        pass
-    elif is_float:
-        has_inf = np.isinf(X).any()
-        has_nan = np.isnan(X).any()
-        if has_inf or not allow_nan and has_nan:
-            if not allow_nan and has_nan:
+    if is_float:
+        with np.errstate(over="ignore"):
+            first_pass_isfinite = np.isfinite(np.sum(X))
+        if first_pass_isfinite:
+            return
+        # Cython implementation doesn't support FP16 or complex numbers
+        use_cython = X.data.contiguous and X.dtype.type in {np.float32, np.float64}
+        if use_cython:
+            out = cy_isfinite(X.reshape(-1), allow_nan=allow_nan)
+            has_nan_error = False if allow_nan else out == FiniteStatus.has_nan
+            has_inf = out == FiniteStatus.has_infinite
+        else:
+            has_inf = np.isinf(X).any()
+            has_nan_error = False if allow_nan else np.isnan(X).any()
+        if has_inf or has_nan_error:
+            if has_nan_error:
                 type_err = "NaN"
             else:
                 msg_dtype = msg_dtype if msg_dtype is not None else X.dtype
                 type_err = f"infinity or a value too large for {msg_dtype!r}"
             padded_input_name = input_name + " " if input_name else ""
             msg_err = f"Input {padded_input_name}contains {type_err}."
-            if not allow_nan and estimator_name and input_name == "X" and has_nan:
+            if estimator_name and input_name == "X" and has_nan_error:
                 # Improve the error message on how to handle missing values in
                 # scikit-learn.
                 msg_err += (