Add tests specified in docstring

Author: paddyroddy

src/array_api_extra/_delegation.py

@@ -114,7 +114,7 @@ def isclose(
 
 
 def nan_to_num(
-    x: Array,
+    x: Array | float | complex,
     /,
     *,
     fill_value: int | float | complex = 0.0,
@@ -136,7 +136,7 @@ def nan_to_num(
 
     Parameters
     ----------
-    x : array
+    x : array, float, complex
         Input data.
     fill_value : int, float, complex, optional
         Value to be used to fill NaN values. If no value is passed
@@ -173,21 +173,20 @@ def nan_to_num(
              0.00000000e+000 +0.00000000e+000j,
              0.00000000e+000 +1.79769313e+308j])
     """
-    if x.ndim == 0:
-        msg = "x must be an array."
-        raise TypeError(msg)
-
     xp = array_namespace(x) if xp is None else xp
 
+    # for scalars we want to output an array
+    y = xp.asarray(x)
+
     if (
         is_cupy_namespace(xp)
         or is_jax_namespace(xp)
         or is_numpy_namespace(xp)
         or is_torch_namespace(xp)
     ):
-        return xp.nan_to_num(x, nan=fill_value)
+        return xp.nan_to_num(y, nan=fill_value)
 
-    return _funcs.nan_to_num(x, fill_value=fill_value, xp=xp)
+    return _funcs.nan_to_num(y, fill_value=fill_value, xp=xp)
 
 
 def one_hot(

src/array_api_extra/_lib/_funcs.py

@@ -760,16 +760,16 @@ def perform_replacements(
         finfo = xp.finfo(x.dtype)
         idx_posinf = xp.isinf(x) & ~xp.signbit(x)
         idx_neginf = xp.isinf(x) & xp.signbit(x)
-        x = xp.where(idx_posinf, x, finfo.max)
-        return xp.where(idx_neginf, x, finfo.min)
+        x = xp.where(idx_posinf, finfo.max, x)
+        return xp.where(idx_neginf, finfo.min, x)
 
     if xp.isdtype(x.dtype, "complex floating"):
         return perform_replacements(
-            x,
+            xp.real(x),
             fill_value,
             xp,
         ) + 1j * perform_replacements(
-            x,
+            xp.imag(x),
             fill_value,
             xp,
         )

tests/conftest.py

@@ -232,3 +232,10 @@ def device(
     if library == Backend.TORCH_GPU:
         return xp.device("cpu")
     return get_device(xp.empty(0))
+
+@pytest.fixture
+def infinity(library: Backend) -> float:
+    """Retrieve the positive infinity value for the given backend."""
+    if library in (Backend.TORCH, Backend.TORCH_GPU):
+        return 3.4028235e+38
+    return 1.7976931348623157e+308

tests/test_funcs.py

@@ -943,6 +943,58 @@ def test_xp(self, xp: ModuleType):
         xp_assert_equal(kron(a, b, xp=xp), k)
 
 
+class TestNumToNan:
+    def test_bool(self, xp: ModuleType) -> None:
+        a = xp.asarray([True])
+        xp_assert_equal(nan_to_num(a), a)
+
+    def test_scalar_pos_inf(self, xp: ModuleType, infinity: float) -> None:
+        a = xp.inf
+        xp_assert_equal(nan_to_num(a, xp=xp), xp.asarray(infinity))
+
+    def test_scalar_neg_inf(self, xp: ModuleType, infinity: float) -> None:
+        a = -xp.inf
+        xp_assert_equal(nan_to_num(a, xp=xp), -xp.asarray(infinity))
+
+    def test_scalar_nan(self, xp: ModuleType) -> None:
+        a = xp.nan
+        xp_assert_equal(nan_to_num(a, xp=xp), xp.asarray(0.0))
+
+    def test_real(self, xp: ModuleType, infinity: float) -> None:
+        a = xp.asarray([xp.inf, -xp.inf, xp.nan, -128, 128])
+        xp_assert_equal(
+            nan_to_num(a),
+            xp.asarray(
+                [
+                    infinity,
+                    -infinity,
+                    0.0,
+                    -128,
+                    128,
+                ]
+            ),
+        )
+
+    def test_complex(self, xp: ModuleType, infinity: float) -> None:
+        a = xp.asarray(
+            [
+                complex(xp.inf, xp.nan),
+                xp.nan,
+                complex(xp.nan, xp.inf),
+            ]
+        )
+        xp_assert_equal(
+            nan_to_num(a),
+            xp.asarray(
+                [
+                    infinity + 0j,
+                    0 + 0j,
+                    0 + 1j * infinity
+                ]
+            ),
+        )
+
+
 class TestNUnique:
     def test_simple(self, xp: ModuleType):
         a = xp.asarray([[1, 1], [0, 2], [2, 2]])