Handle case of complex fill val

Author: paddyroddy

src/array_api_extra/_lib/_funcs.py

@@ -762,14 +762,14 @@ def perform_replacements(  # numpydoc ignore=PR01,RT01
         x = xp.where(idx_posinf, finfo.max, x)
         return xp.where(idx_neginf, finfo.min, x)
 
-    if xp.isdtype(x.dtype, "complex floating"):
+    if isinstance(fill_value, complex) or xp.isdtype(x.dtype, "complex floating"):
         return perform_replacements(
             xp.real(x),
-            fill_value,
+            fill_value.real,
             xp,
         ) + 1j * perform_replacements(
             xp.imag(x),
-            fill_value,
+            fill_value.imag,
             xp,
         )
 

tests/test_funcs.py

@@ -988,11 +988,25 @@ def test_complex(self, xp: ModuleType, infinity: float) -> None:
             xp.asarray([infinity + 0j, 0 + 0j, 0 + 1j * infinity]),
         )
 
-    def test_fill_value(self, xp: ModuleType) -> None:
+    @pytest.mark.parametrize("fill_value", [3, 3.0, 3 + 0j])
+    @pytest.mark.parametrize(
+        "output",
+        [
+            [1.0, 2.0, 3.0, 4.0],
+            [1.0, 2.0, 3.0, 4.0],
+            [1.0 + 0.j, 2.0 + 0.j, 3.0 + 0.j, 4.0 + 0.j]
+        ],
+    )
+    def test_fill_value(
+        self,
+        xp: ModuleType,
+        fill_value: float,
+        output: float,
+    ) -> None:
         a = xp.asarray([1, 2, np.nan, 4])
         xp_assert_equal(
-            nan_to_num(a, fill_value=3, xp=xp),
-            xp.asarray([1.0, 2.0, 3.0, 4.0]),
+            nan_to_num(a, fill_value=fill_value, xp=xp),
+            xp.asarray(output),
         )
 
     def test_empty_array(self, xp: ModuleType) -> None: