tnp.put and test_put

Author: honno

torch_np/_funcs_impl.py

@@ -913,6 +913,28 @@ def take_along_axis(arr: ArrayLike, indices: ArrayLike, axis):
     return torch.take_along_dim(arr, indices, axis)
 
 
+def put(a: ArrayLike, ind: Sequence[ArrayLike], v: ArrayLike, mode="raise"):
+    if mode != "raise":
+        raise NotImplementedError(f"{mode=}")
+
+    index = torch.concat(ind)
+    index[index < 0] += a.numel()  # normalise negative indices
+    index_u, index_c = torch.unique(index, return_counts=True)
+    duplicated_indices = index_u[index_c > 1]
+    if duplicated_indices.numel() > 0:
+        raise NotImplementedError(
+            "duplicated indices are not supported. duplicated indices: "
+            f"{list(duplicated_indices)}"
+        )
+    source = v
+    if source.numel() < index.numel():
+        source = torch.broadcast_to(source, index.size())
+    # Note Tensor.put_ acts in-place, while Tensor.put (no trailing underscore)
+    # is seemingly out-of-place.
+    a.put_(index, source)
+    return None
+
+
 def put_along_axis(arr: ArrayLike, indices: ArrayLike, values: ArrayLike, axis):
     (arr,), axis = _util.axis_none_ravel(arr, axis=axis)
     axis = _util.normalize_axis_index(axis, arr.ndim)

torch_np/tests/test_xps.py

@@ -10,30 +10,35 @@
 
 pytest.importorskip("hypothesis")
 
-from hypothesis import given
+import numpy as np
+import torch
+from hypothesis import given, note
 from hypothesis import strategies as st
 from hypothesis.errors import HypothesisWarning
+from hypothesis.extra import numpy as nps
 from hypothesis.extra.array_api import make_strategies_namespace
 
-import torch_np as np
+import torch_np as tnp
+from torch_np._dtypes import sctypes
+from torch_np.testing import assert_array_equal
 
 __all__ = ["xps"]
 
 with warnings.catch_warnings():
     warnings.filterwarnings("ignore", category=HypothesisWarning)
-    np.bool = np.bool_
-    xps = make_strategies_namespace(np, api_version="2022.12")
+    tnp.bool = tnp.bool_
+    xps = make_strategies_namespace(tnp, api_version="2022.12")
 
 
-default_dtypes = [np.bool, np.int64, np.float64, np.complex128]
+default_dtypes = [tnp.bool, tnp.int64, tnp.float64, tnp.complex128]
 kind_to_strat = {
     "b": xps.boolean_dtypes(),
     "i": xps.integer_dtypes(),
     "u": xps.unsigned_integer_dtypes(sizes=8),
     "f": xps.floating_dtypes(),
     "c": xps.complex_dtypes(),
 }
-scalar_dtype_strat = st.one_of(kind_to_strat.values()).map(np.dtype)
+scalar_dtype_strat = st.one_of(kind_to_strat.values()).map(tnp.dtype)
 
 
 @pytest.mark.skip(reason="flaky")
@@ -55,14 +60,14 @@ def test_full(shape, data):
         else:
             values_dtypes_strat = st.just(_dtype)
         values_strat = values_dtypes_strat.flatmap(
-            lambda d: values_strat.map(lambda v: np.asarray(v, dtype=d))
+            lambda d: values_strat.map(lambda v: tnp.asarray(v, dtype=d))
         )
     fill_value = data.draw(values_strat, label="fill_value")
-    out = np.full(shape, fill_value, **kw)
+    out = tnp.full(shape, fill_value, **kw)
     assert out.dtype == _dtype
     assert out.shape == shape
     if cmath.isnan(fill_value):
-        assert np.isnan(out).all()
+        assert tnp.isnan(out).all()
     else:
         assert (out == fill_value).all()
 
@@ -89,3 +94,48 @@ 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(
+    np_x=nps.arrays(
+        # We specifically use namespaced dtypes to prevent non-native byte-order issues
+        dtype=scalar_dtype_strat.map(lambda d: getattr(np, d.name)),
+        shape=nps.array_shapes(),
+    ),
+    data=st.data(),
+)
+def test_put(np_x, data):
+    # We cast arrays from torch_np.asarray as currently it doesn't carry over
+    # dtypes. XXX: Remove the below sanity check and subsequent casting when
+    # this is fixed.
+    assert tnp.asarray(np.zeros(5, dtype=np.int16)).dtype != tnp.int16
+
+    tnp_x = tnp.asarray(np_x.copy()).astype(np_x.dtype.name)
+
+    result_shapes = st.shared(nps.array_shapes())
+    ind = data.draw(
+        nps.integer_array_indices(np_x.shape, result_shape=result_shapes), label="ind"
+    )
+    v = data.draw(nps.arrays(dtype=np_x.dtype, shape=result_shapes), label="v")
+
+    tnp_x_copy = tnp_x.copy()
+    np.put(np_x, ind, v)
+    note(f"(after put) {np_x=}")
+    assert_array_equal(tnp_x, tnp_x_copy)  # sanity check
+
+    note(f"{tnp_x=}")
+    tnp_ind = []
+    for np_indices in ind:
+        tnp_indices = tnp.asarray(np_indices).astype(np_indices.dtype.name)
+        tnp_ind.append(tnp_indices)
+    tnp_ind = tuple(tnp_ind)
+    note(f"{tnp_ind=}")
+    tnp_v = tnp.asarray(v.copy()).astype(v.dtype.name)
+    note(f"{tnp_v=}")
+    try:
+        tnp.put(tnp_x, tnp_ind, tnp_v)
+    except NotImplementedError:
+        return
+    note(f"(after put) {tnp_x=}")
+
+    assert_array_equal(tnp_x, tnp.asarray(np_x).astype(tnp_x.dtype))