Smoke test for advance integer indexing

Author: honno

torch_np/_helpers.py

@@ -74,9 +74,23 @@ def ndarrays_to_tensors(*inputs):
     """Convert all ndarrays from `inputs` to tensors. (other things are intact)"""
     from ._ndarray import asarray, ndarray
 
-    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

@@ -410,11 +410,9 @@ def clip(self, min, max, out=None):
     )
 
     ### 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

@@ -312,6 +312,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

@@ -1,10 +1,13 @@
 import itertools
 
 import pytest
-from pytest import raises as assert_raises
 
 # import numpy as np
+import torch
+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
 
 

torch_np/tests/test_ndarray_methods.py

@@ -58,3 +58,27 @@ def test_full(shape, data):
         assert np.isnan(out).all()
     else:
         assert (out == fill_value).all()
+
+
+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