Test for full()

Author: honno

torch_np/tests/test_stuff.py

@@ -9,13 +9,14 @@
 from hypothesis.errors import HypothesisWarning
 from hypothesis.extra.array_api import make_strategies_namespace
 
-import torch_np
+import torch_np as np
 
 __all__ = ["xps"]
 
 with warnings.catch_warnings():
     warnings.filterwarnings("ignore", category=HypothesisWarning)
-    xps = make_strategies_namespace(torch_np, api_version="2021.12")
+    np.bool = np.bool_
+    xps = make_strategies_namespace(np, api_version="2021.12")
 
 
 def integer_array_indices(shape, result_shape) -> st.SearchStrategy[tuple]:
@@ -40,3 +41,29 @@ def test_integer_indexing(x, data):
     idx = data.draw(integer_array_indices(x.shape, result_shape), label="idx")
     result = x[idx]
     assert result.shape == result_shape
+
+
+@given(shape=xps.array_shapes(), data=st.data())
+def test_full(shape, data):
+    if data.draw(st.booleans(), label="pass kwargs?"):
+        kw = {}
+    else:
+        dtype = data.draw(st.none() | xps.scalar_dtypes(), label="dtype")
+        kw = {"dtype": dtype}
+    _dtype = kw.get("dtype", None) or data.draw(
+        st.sampled_from([np.bool, np.int64, np.float64]), label="_dtype"
+    )
+    fill_value = data.draw(xps.from_dtype(_dtype), label="fill_value")
+    out = np.full(shape, fill_value, **kw)
+    if kw.get("dtype", None) is None:
+        if isinstance(fill_value, bool):
+            assert out.dtype == np.bool
+        elif isinstance(fill_value, int):
+            assert out.dtype == np.int64
+        else:
+            assert isinstance(fill_value, float)  # sanity check
+            assert out.dtype == np.float64
+    else:
+        assert out.dtype == kw["dtype"]
+    assert out.shape == shape
+    assert (out == fill_value).all()