Port indexing tests

torch_np/__init__.py

@@ -1,6 +1,7 @@
 from ._wrapper import *  # isort: skip  # XXX: currently this prevents circular imports
 from . import random
 from ._binary_ufuncs import *
+from ._detail._index_tricks import *
 from ._detail._util import AxisError, UFuncTypeError
 from ._dtypes import *
 from ._getlimits import errstate, finfo, iinfo
@@ -15,3 +16,6 @@
 inf = float("inf")
 nan = float("nan")
 from math import pi  # isort: skip
+
+False_ = asarray(False, bool_)
+True_ = asarray(True, bool_)

torch_np/_detail/_index_tricks.py

@@ -0,0 +1,26 @@
+"""
+Vendored objects from numpy.lib.index_tricks
+"""
+__all__ = ["index_exp", "s_"]
+
+
+class IndexExpression:
+    """
+    Written by Konrad Hinsen <[email protected]>
+    last revision: 1999-7-23
+
+    Cosmetic changes by T. Oliphant 2001
+    """
+
+    def __init__(self, maketuple):
+        self.maketuple = maketuple
+
+    def __getitem__(self, item):
+        if self.maketuple and not isinstance(item, tuple):
+            return (item,)
+        else:
+            return item
+
+
+index_exp = IndexExpression(maketuple=True)
+s_ = IndexExpression(maketuple=False)

torch_np/_detail/implementations.py

@@ -1,5 +1,6 @@
 import torch
 
+from .. import _helpers
 from . import _dtypes_impl, _util
 
 # ### equality, equivalence, allclose ###
@@ -503,8 +504,10 @@ def reshape(tensor, *shape, order="C"):
     if order != "C":
         raise NotImplementedError
     newshape = shape[0] if len(shape) == 1 else shape
+    # convert any tnp.ndarray inputs into tensors before passing to torch.Tensor.reshape
+    t_newshape = _helpers.ndarrays_to_tensors(newshape)
     # if sh = (1, 2, 3), numpy allows both .reshape(sh) and .reshape(*sh)
-    result = tensor.reshape(newshape)
+    result = tensor.reshape(t_newshape)
     return result
 
 

torch_np/_dtypes.py

@@ -161,23 +161,23 @@ class bool_(generic):
 
 
 # name aliases : FIXME (OS, bitness)
-intp = int64
-int_ = int64
-intc = int32
-
-byte = int8
-short = int16
-longlong = int64  # XXX: is this correct?
-
-ubyte = uint8
-
-half = float16
-single = float32
-double = float64
-float_ = float64
-
-csingle = complex64
-cdouble = complex128
+_name_aliases = {
+    "intp": int64,
+    "int_": int64,
+    "intc": int32,
+    "byte": int8,
+    "short": int16,
+    "longlong": int64,  # XXX: is this correct?
+    "ubyte": uint8,
+    "half": float16,
+    "single": float32,
+    "double": float64,
+    "float_": float64,
+    "csingle": complex64,
+    "cdouble": complex128,
+}
+for name, obj in _name_aliases.items():
+    globals()[name] = obj
 
 
 # Replicate this NumPy-defined way of grouping scalar types,
@@ -232,6 +232,8 @@ def sctype_from_string(s):
     """Normalize a string value: a type 'name' or a typecode or a width alias."""
     if s in _names:
         return _names[s]
+    if s in _name_aliases.keys():
+        return _name_aliases[s]
     if s in _typecodes:
         return _typecodes[s]
     if s in _aliases:

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

@@ -381,14 +381,24 @@ def repeat(self, repeats, axis=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
-        )
+    @staticmethod
+    def _upcast_int_indices(index):
+        if isinstance(index, torch.Tensor):
+            if index.dtype in (torch.int8, torch.int16, torch.int32, torch.uint8):
+                return index.to(torch.int64)
+        elif isinstance(index, tuple):
+            return tuple(ndarray._upcast_int_indices(i) for i in index)
+        return index
+
+    def __getitem__(self, index):
+        index = _helpers.ndarrays_to_tensors(index)
+        index = ndarray._upcast_int_indices(index)
+        return ndarray._from_tensor_and_base(self._tensor.__getitem__(index), self)
 
     def __setitem__(self, index, value):
-        value = asarray(value).get()
+        index = _helpers.ndarrays_to_tensors(index)
+        index = ndarray._upcast_int_indices(index)
+        value = _helpers.ndarrays_to_tensors(value)
         return self._tensor.__setitem__(index, value)
 
     ### sorting ###

torch_np/_wrapper.py

@@ -291,6 +291,8 @@ def empty_like(prototype, dtype=None, order="K", subok=False, shape=None):
 @_decorators.dtype_to_torch
 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
     fill_value = asarray(fill_value).get()

torch_np/testing/__init__.py

@@ -1,4 +1,6 @@
 from .utils import (
+    HAS_REFCOUNT,
+    IS_WASM,
     _gen_alignment_data,
     assert_,
     assert_allclose,

torch_np/testing/utils.py

@@ -17,8 +17,6 @@
 from tempfile import mkdtemp, mkstemp
 from warnings import WarningMessage
 
-from pytest import raises as assert_raises
-
 import torch_np as np
 from torch_np import arange, array
 from torch_np import asarray as asanyarray
@@ -31,8 +29,8 @@
     "assert_array_equal",
     "assert_array_less",
     "assert_string_equal",
+    "assert_",
     "assert_array_almost_equal",
-    "assert_raises",
     "build_err_msg",
     "decorate_methods",
     "print_assert_equal",