ENH: add where, searchsorted, inner and outer, diag_indices{_from}

Author: ev-br

autogen/numpy_api_dump.py

@@ -216,22 +216,6 @@ def deprecate_with_doc(msg):
     raise NotImplementedError
 
 
-def diag_indices(n, ndim=2):
-    raise NotImplementedError
-
-
-def diag_indices_from(arr):
-    raise NotImplementedError
-
-
-def diagflat(v, k=0):
-    raise NotImplementedError
-
-
-def diagonal(a, offset=0, axis1=0, axis2=1):
-    raise NotImplementedError
-
-
 def digitize(x, bins, right=False):
     raise NotImplementedError
 
@@ -397,10 +381,6 @@ def indices(dimensions, dtype=int, sparse=False):
     raise NotImplementedError
 
 
-def inner(a, b, /):
-    raise NotImplementedError
-
-
 def insert(arr, obj, values, axis=None):
     raise NotImplementedError
 
@@ -417,10 +397,6 @@ def is_busday(dates, weekmask="1111100", holidays=None, busdaycal=None, out=None
     raise NotImplementedError
 
 
-def isclose(a, b, rtol=1e-05, atol=1e-08, equal_nan=False):
-    raise NotImplementedError
-
-
 def isfortran(a):
     raise NotImplementedError
 
@@ -437,14 +413,6 @@ def issctype(rep):
     raise NotImplementedError
 
 
-def issubclass_(arg1, arg2):
-    raise NotImplementedError
-
-
-def issubdtype(arg1, arg2):
-    raise NotImplementedError
-
-
 def issubsctype(arg1, arg2):
     raise NotImplementedError
 
@@ -624,10 +592,6 @@ def obj2sctype(rep, default=None):
     raise NotImplementedError
 
 
-def outer(a, b, out=None):
-    raise NotImplementedError
-
-
 def packbits(a, /, axis=None, bitorder="big"):
     raise NotImplementedError
 
@@ -696,10 +660,6 @@ def putmask(a, mask, values):
     raise NotImplementedError
 
 
-def ravel(a, order="C"):
-    raise NotImplementedError
-
-
 def recfromcsv(fname, **kwargs):
     raise NotImplementedError
 
@@ -754,10 +714,6 @@ def sctype2char(sctype):
     raise NotImplementedError
 
 
-def searchsorted(a, v, side="left", sorter=None):
-    raise NotImplementedError
-
-
 def select(condlist, choicelist, default=0):
     raise NotImplementedError
 
@@ -810,23 +766,10 @@ def shares_memory(a, b, max_work=None):
 def show():
     raise NotImplementedError
 
-
-def sinc(x):
-    raise NotImplementedError
-
-
-def sometrue(*args, **kwargs):
-    raise NotImplementedError
-
-
 def sort_complex(a):
     raise NotImplementedError
 
 
-def swapaxes(a, axis1, axis2):
-    raise NotImplementedError
-
-
 def take(a, indices, axis=None, out=None, mode="raise"):
     raise NotImplementedError
 
@@ -879,9 +822,5 @@ def vdot(a, b, /):
     raise NotImplementedError
 
 
-def where(condition, x, y, /):
-    raise NotImplementedError
-
-
 def who(vardict=None):
     raise NotImplementedError

torch_np/_detail/implementations.py

@@ -127,6 +127,22 @@ def triu_indices(n, k=0, m=None):
     return result
 
 
+def diag_indices(n, ndim=2):
+    idx = torch.arange(n)
+    return (idx,)*ndim
+
+
+def diag_indices_from(tensor):
+    if not tensor.ndim >= 2:
+        raise ValueError("input array must be at least 2-d")
+    # For more than d=2, the strided formula is only valid for arrays with
+    # all dimensions equal, so we check first.
+    s = tensor.shape
+    if any(s[1:] != s[:-1]):
+        raise ValueError("All dimensions of input must be of equal length")
+    return diag_indices(s[0], tensor.ndim)
+
+
 # ### splits ###
 
 
@@ -589,3 +605,18 @@ def argsort(tensor, axis=-1, kind=None, order=None):
     tensor, axis, stable = _sort_helper(tensor, axis, kind, order)
     result = torch.argsort(tensor, dim=axis, stable=stable)
     return result
+
+
+# ### logic and selection ###
+
+
+def where(condition, x, y):
+    selector = (x is None) == (y is None)
+    if not selector:
+        raise ValueError("either both or neither of x and y should be given")
+
+    if x is None and y is None:
+        result = torch.where(condition)
+    else:
+        result = torch.where(condition, x, y)
+    return result

torch_np/_ndarray.py

@@ -395,6 +395,10 @@ def repeat(self, repeats, axis=None):
         axis_none_ravel_wrapper(dtype_to_torch(_reductions.cumsum))
     )
 
+    def diagonal(self, offset=0, axis1=0, axis2=1):
+        result = torch.diagonal(self._tensor, offset, axis1, axis2)
+        return asarray(result)
+
     ### indexing ###
     @staticmethod
     def _upcast_int_indices(index):
@@ -427,6 +431,11 @@ def argsort(self, axis=-1, kind=None, order=None):
         result = _impl.argsort(self.tensor, axis, kind, order)
         return asarray(result)
 
+    def searchsorted(self, v, side="left", sorter=None):
+        v_t, sorter_t = _helpers.to_tensors_or_none(v, sorter)
+        result = torch.searchsorted(self._tensor, v_t, side=side, sorter=sorter_t)
+        return asarray(result)
+
 
 # This is the ideally the only place which talks to ndarray directly.
 # The rest goes through asarray (preferred) or array.

torch_np/_wrapper.py

@@ -372,6 +372,28 @@ def diag(v, k=0):
     return asarray(result)
 
 
+def diagonal(a, offset=0, axis1=0, axis2=1):
+    arr = asarray(a)
+    return arr.diagonal(offset, axis1, axis2)
+
+
+def diagflat(v, k=0):
+    tensor = asarray(v).get()
+    result = torch.diagflat(tensor, k)
+    return result
+
+
+def diag_indices(n, ndim=2):
+    result = _impl.diag_indices(n, ndim)
+    return tuple(asarray(x) for x in result)
+
+
+def diag_indices_from(arr):
+    tensor = asarray(arr).get()
+    result = _impl.diag_indices_from(tensor)
+    return tuple(asarray(x) for x in result)
+
+
 ###### misc/unordered
 
 
@@ -446,17 +468,19 @@ def bincount(x, /, weights=None, minlength=0):
     return asarray(result)
 
 
-# YYY: pattern: sequence of arrays
 def where(condition, x=None, y=None, /):
-    selector = (x is None) == (y is None)
-    if not selector:
-        raise ValueError("either both or neither of x and y should be given")
-    condition = asarray(condition).get()
-    if x is None and y is None:
-        return tuple(asarray(_) for _ in torch.where(condition))
-    x = asarray(condition).get()
-    y = asarray(condition).get()
-    return asarray(torch.where(condition, x, y))
+    cond_t, x_t, y_t = _helpers.to_tensors_or_none(condition, x, y)
+    result = _impl.where(cond_t, x_t, y_t)
+    if isinstance(result, tuple):
+        # single-argument where(condition)
+        return tuple(asarray(x) for x in result)
+    else:
+        return asarray(result)
+
+
+def searchsorted(a, v, side="left", sorter=None):
+    arr = asarray(a)
+    return arr.searchsorted(v, side=side, sorter=sorter)
 
 
 ###### module-level queries of object properties
@@ -847,6 +871,18 @@ def median(a, axis=None, out=None, overwrite_input=False, keepdims=False):
     )
 
 
+def inner(a, b, /):
+    a_t, b_t = _helpers.to_tensors(a, b)
+    result = torch.inner(a_t, b_t)
+    return asarray(result)
+
+
+def outer(a, b, out=None):
+    a_t, b_t = _helpers.to_tensors(a, b)
+    result = torch.outer(a_t, b_t)
+    return _helpers.result_or_out(result, out)
+
+
 @asarray_replacer()
 def nanmean(a, axis=None, dtype=None, out=None, keepdims=NoValue, *, where=NoValue):
     if where is not NoValue:

torch_np/tests/numpy_tests/core/test_numeric.py

@@ -125,7 +125,6 @@ def test_cumproduct(self):
         A = [[1, 2, 3], [4, 5, 6]]
         assert_(np.all(np.cumproduct(A) == np.array([1, 2, 6, 24, 120, 720])))
 
-    @pytest.mark.xfail(reason="TODO implement diagonal(...)")
     def test_diagonal(self):
         a = [[0, 1, 2, 3],
              [4, 5, 6, 7],