Merge pull request #272 from hameerabbasi/fix-nan-propagation-immutable

Fix `test_nan_propagation` for immutable arrays

Author: asmeurer

array_api_tests/test_special_cases.py

@@ -23,7 +23,7 @@
 from warnings import warn
 
 import pytest
-from hypothesis import given, note, settings
+from hypothesis import given, note, settings, assume
 from hypothesis import strategies as st
 
 from array_api_tests.typing import Array, DataType
@@ -1331,10 +1331,11 @@ def test_empty_arrays(func_name, expected):  # TODO: parse docstrings to get exp
 )
 def test_nan_propagation(func_name, x, data):
     func = getattr(xp, func_name)
-    set_idx = data.draw(
-        xps.indices(x.shape, max_dims=0, allow_ellipsis=False), label="set idx"
+    nan_positions = data.draw(
+        hh.arrays(dtype=hh.bool_dtype, shape=x.shape), label="nan_positions"
     )
-    x[set_idx] = float("nan")
+    assume(xp.any(nan_positions))
+    x = xp.where(nan_positions, xp.asarray(float("nan")), x)
     note(f"{x=}")
 
     out = func(x)