Smoke test for advance integer indexing

Author: honno

torch_np/_helpers.py

@@ -56,9 +56,23 @@ def result_or_out(result_tensor, out_array=None):
 
 def ndarrays_to_tensors(*inputs):
     """Convert all ndarrays from `inputs` to tensors. (other things are intact)"""
-    return tuple(
-        [value.get() if isinstance(value, ndarray) else value for value in inputs]
-    )
+    if len(inputs) == 0:
+        return ValueError()
+    elif len(inputs) == 1:
+        input_ = inputs[0]
+        if isinstance(input_, ndarray):
+            return input_.get()
+        elif isinstance(input_, tuple):
+            result = []
+            for sub_input in input_:
+                sub_result = ndarrays_to_tensors(sub_input)
+                result.append(sub_result)
+            return tuple(result)
+        else:
+            return input_
+    else:
+        assert isinstance(inputs, tuple)  # sanity check
+        return ndarrays_to_tensors(inputs)
 
 
 def to_tensors(*inputs):

torch_np/_ndarray.py

@@ -282,11 +282,9 @@ def nonzero(self):
     std = emulate_out_arg(axis_keepdims_wrapper(dtype_to_torch(_reductions.std)))
 
     ### indexing ###
-    def __getitem__(self, *args, **kwds):
-        t_args = _helpers.ndarrays_to_tensors(*args)
-        return ndarray._from_tensor_and_base(
-            self._tensor.__getitem__(*t_args, **kwds), self
-        )
+    def __getitem__(self, index):
+        t_index = _helpers.ndarrays_to_tensors(index)
+        return ndarray._from_tensor_and_base(self._tensor.__getitem__(t_index), self)
 
     def __setitem__(self, index, value):
         value = asarray(value).get()

torch_np/_wrapper.py

@@ -233,6 +233,8 @@ def empty_like(prototype, dtype=None, order="K", subok=False, shape=None):
 
 def full(shape, fill_value, dtype=None, order="C", *, like=None):
     _util.subok_not_ok(like)
+    if isinstance(shape, int):
+        shape = (shape,)
     if order != "C":
         raise NotImplementedError
     if isinstance(fill_value, ndarray):

torch_np/tests/__init__.py

@@ -0,0 +1,12 @@
+import warnings
+
+from hypothesis.errors import HypothesisWarning
+from hypothesis.extra.array_api import make_strategies_namespace
+
+import torch_np
+
+__all__ = ["xps"]
+
+with warnings.catch_warnings():
+    warnings.filterwarnings("ignore", category=HypothesisWarning)
+    xps = make_strategies_namespace(torch_np, api_version="2021.12")

torch_np/tests/test_ndarray_methods.py

@@ -1,12 +1,19 @@
 import itertools
 
 import pytest
-from pytest import raises as assert_raises
 
 # import numpy as np
+import torch
+from hypothesis import given
+from hypothesis import strategies as st
+from pytest import raises as assert_raises
+
 import torch_np as np
+from torch_np._ndarray import ndarray
 from torch_np.testing import assert_equal
 
+from . import xps
+
 
 class TestIndexing:
     def test_indexing_simple(self):
@@ -23,6 +30,30 @@ def test_setitem(self):
         assert_equal(a, [[8, 2, 3], [4, 5, 6]])
 
 
+def integer_array_indices(shape, result_shape) -> st.SearchStrategy[tuple]:
+    # See hypothesis.extra.numpy.integer_array_indices()
+    # n.b. result_shape only accepts a shape, as opposed to only accepting a strategy
+    def array_for(index_shape, size):
+        return xps.arrays(
+            dtype=xps.integer_dtypes(),
+            shape=index_shape,
+            elements=st.integers(-size, size - 1),
+        )
+
+    return st.tuples(*(array_for(result_shape, size) for size in shape))
+
+
+@given(
+    x=xps.arrays(dtype=xps.integer_dtypes(), shape=xps.array_shapes()),
+    data=st.data(),
+)
+def test_integer_indexing(x, data):
+    result_shape = data.draw(xps.array_shapes(), label="result_shape")
+    idx = data.draw(integer_array_indices(x.shape, result_shape), label="idx")
+    result = x[idx]
+    assert result.shape == result_shape
+
+
 class TestReshape:
     def test_reshape_function(self):
         arr = [[1, 2, 3], [4, 5, 6], [7, 8, 9], [10, 11, 12]]