feat: add finite/infinite/nan checks

decent_array/interoperability/__init__.py

@@ -50,6 +50,9 @@
     add,
     divide,
     floor_divide,
+    isfinite,
+    isinf,
+    isnan,
     multiply,
     negative,
     positive,
@@ -111,6 +114,9 @@
     "get_seed",
     "greater",
     "greater_equal",
+    "isfinite",
+    "isinf",
+    "isnan",
     "less",
     "less_equal",
     "matmul",

decent_array/interoperability/_abstracts/backend.py

@@ -270,6 +270,18 @@ def absolute(self, x: Array) -> Array:
     def sqrt(self, x: Array) -> Array:
         """Element-wise square root."""
 
+    @abstractmethod
+    def isfinite(self, x: Array) -> Array:
+        """Element-wise test for finite values."""
+
+    @abstractmethod
+    def isinf(self, x: Array) -> Array:
+        """Element-wise test for infinite values."""
+
+    @abstractmethod
+    def isnan(self, x: Array) -> Array:
+        """Element-wise test for NaN values."""
+
     # Comparisons — both operands may be Array or scalar.
 
     @abstractmethod

decent_array/interoperability/_iop/math.py

@@ -147,3 +147,24 @@ def sqrt(x: Array) -> Array:
     if _BACKEND_INSTANCE is None:
         raise no_backend_error
     return _BACKEND_INSTANCE.sqrt(x)
+
+
+def isfinite(x: Array) -> Array:
+    """Element-wise test for finite values."""
+    if _BACKEND_INSTANCE is None:
+        raise no_backend_error
+    return _BACKEND_INSTANCE.isfinite(x)
+
+
+def isinf(x: Array) -> Array:
+    """Element-wise test for infinite values."""
+    if _BACKEND_INSTANCE is None:
+        raise no_backend_error
+    return _BACKEND_INSTANCE.isinf(x)
+
+
+def isnan(x: Array) -> Array:
+    """Element-wise test for NaN values."""
+    if _BACKEND_INSTANCE is None:
+        raise no_backend_error
+    return _BACKEND_INSTANCE.isnan(x)

decent_array/interoperability/_jax/jax_backend.py

@@ -255,6 +255,18 @@ def absolute(self, x: Array) -> Array:
     def sqrt(self, x: Array) -> Array:
         return Array(jnp.sqrt(x.value))
 
+    def isfinite(self, x: Array) -> Array:
+        """Element-wise test for finite values."""
+        return Array(jnp.isfinite(x.value))
+
+    def isinf(self, x: Array) -> Array:
+        """Element-wise test for infinite values."""
+        return Array(jnp.isinf(x.value))
+
+    def isnan(self, x: Array) -> Array:
+        """Element-wise test for NaN values."""
+        return Array(jnp.isnan(x.value))
+
     # Comparisons
 
     def equal(self, x1: int | float | complex | Array, x2: int | float | complex | Array) -> Array:

decent_array/interoperability/_numpy/numpy_backend.py

@@ -267,6 +267,18 @@ def absolute(self, x: Array) -> Array:
     def sqrt(self, x: Array) -> Array:
         return Array(np.sqrt(x.value))
 
+    def isfinite(self, x: Array) -> Array:
+        """Element-wise test for finite values."""
+        return Array(np.isfinite(x.value))
+
+    def isinf(self, x: Array) -> Array:
+        """Element-wise test for infinite values."""
+        return Array(np.isinf(x.value))
+
+    def isnan(self, x: Array) -> Array:
+        """Element-wise test for NaN values."""
+        return Array(np.isnan(x.value))
+
     # Comparisons
 
     def equal(self, x1: int | float | complex | Array, x2: int | float | complex | Array) -> Array:

decent_array/interoperability/_pytorch/pytorch_backend.py

@@ -268,6 +268,18 @@ def absolute(self, x: Array) -> Array:
     def sqrt(self, x: Array) -> Array:
         return Array(torch.sqrt(x.value))
 
+    def isfinite(self, x: Array) -> Array:
+        """Element-wise test for finite values."""
+        return Array(torch.isfinite(x.value))
+
+    def isinf(self, x: Array) -> Array:
+        """Element-wise test for infinite values."""
+        return Array(torch.isinf(x.value))
+
+    def isnan(self, x: Array) -> Array:
+        """Element-wise test for NaN values."""
+        return Array(torch.isnan(x.value))
+
     # Comparisons
 
     def equal(self, x1: int | float | complex | Array, x2: int | float | complex | Array) -> Array:

decent_array/interoperability/_tensorflow/tensorflow_backend.py

@@ -278,6 +278,18 @@ def absolute(self, x: Array) -> Array:
     def sqrt(self, x: Array) -> Array:
         return Array(tf.sqrt(x.value))
 
+    def isfinite(self, x: Array) -> Array:
+        """Element-wise test for finite values."""
+        return Array(tf.math.is_finite(x.value))
+
+    def isinf(self, x: Array) -> Array:
+        """Element-wise test for infinite values."""
+        return Array(tf.math.is_inf(x.value))
+
+    def isnan(self, x: Array) -> Array:
+        """Element-wise test for NaN values."""
+        return Array(tf.math.is_nan(x.value))
+
     # Comparisons
 
     def equal(self, x1: int | float | complex | Array, x2: int | float | complex | Array) -> Array:

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "decent-array"
-version = "0.2.4"
+version = "0.2.5"
 authors = [{name = "Simon Granström"}, {name = "Nicola Bastianello"}]
 maintainers = [{name = "Team Decent"}]
 description = "A library of array operations and linear algebra primitives for interoperability across ML frameworks."

tests/test_iop_functions.py

@@ -478,6 +478,21 @@ def test_sqrt(backend: tuple) -> None:
     np.testing.assert_allclose(_np(iop.sqrt(arr)), [1.0, 2.0, 3.0])
 
 
+def test_isfinite(backend: tuple) -> None:
+    arr = iop.from_numpy(np.array([1.0, np.inf, -np.inf, np.nan], dtype=np.float32))
+    np.testing.assert_array_equal(_np(iop.isfinite(arr)), [True, False, False, False])
+
+
+def test_isinf(backend: tuple) -> None:
+    arr = iop.from_numpy(np.array([1.0, np.inf, -np.inf, np.nan], dtype=np.float32))
+    np.testing.assert_array_equal(_np(iop.isinf(arr)), [False, True, True, False])
+
+
+def test_isnan(backend: tuple) -> None:
+    arr = iop.from_numpy(np.array([1.0, np.inf, -np.inf, np.nan], dtype=np.float32))
+    np.testing.assert_array_equal(_np(iop.isnan(arr)), [False, False, False, True])
+
+
 # Operators --------------------------------------------------------------