ENH: fill_diagonal, trace, indices, vdot

Author: ev-br

autogen/numpy_api_dump.py

@@ -172,10 +172,6 @@ def choose(a, choices, out=None, mode="raise"):
     raise NotImplementedError
 
 
-def clip(a, a_min, a_max, out=None, **kwargs):
-    raise NotImplementedError
-
-
 def common_type(*arrays):
     raise NotImplementedError
 
@@ -244,10 +240,6 @@ def extract(condition, arr):
     raise NotImplementedError
 
 
-def fill_diagonal(a, val, wrap=False):
-    raise NotImplementedError
-
-
 def find_common_type(array_types, scalar_types):
     raise NotImplementedError
 
@@ -377,10 +369,6 @@ def in1d(ar1, ar2, assume_unique=False, invert=False):
     raise NotImplementedError
 
 
-def indices(dimensions, dtype=int, sparse=False):
-    raise NotImplementedError
-
-
 def insert(arr, obj, values, axis=None):
     raise NotImplementedError
 
@@ -766,6 +754,7 @@ def shares_memory(a, b, max_work=None):
 def show():
     raise NotImplementedError
 
+
 def sort_complex(a):
     raise NotImplementedError
 
@@ -778,10 +767,6 @@ def tensordot(a, b, axes=2):
     raise NotImplementedError
 
 
-def trace(a, offset=0, axis1=0, axis2=1, dtype=None, out=None):
-    raise NotImplementedError
-
-
 def trapz(y, x=None, dx=1.0, axis=-1):
     raise NotImplementedError
 
@@ -818,9 +803,5 @@ def unwrap(p, discont=None, axis=-1, *, period=6.283185307179586):
     raise NotImplementedError
 
 
-def vdot(a, b, /):
-    raise NotImplementedError
-
-
 def who(vardict=None):
     raise NotImplementedError

torch_np/_detail/implementations.py

@@ -129,7 +129,7 @@ def triu_indices(n, k=0, m=None):
 
 def diag_indices(n, ndim=2):
     idx = torch.arange(n)
-    return (idx,)*ndim
+    return (idx,) * ndim
 
 
 def diag_indices_from(tensor):
@@ -143,6 +143,34 @@ def diag_indices_from(tensor):
     return diag_indices(s[0], tensor.ndim)
 
 
+def fill_diagonal(tensor, t_val, wrap):
+    # torch.Tensor.fill_diagonal_ only accepts scalars. Thus vendor the numpy source,
+    # https://github.com/numpy/numpy/blob/v1.24.0/numpy/lib/index_tricks.py#L786-L917
+
+    if tensor.ndim < 2:
+        raise ValueError("array must be at least 2-d")
+    end = None
+    if tensor.ndim == 2:
+        # Explicit, fast formula for the common case.  For 2-d arrays, we
+        # accept rectangular ones.
+        step = tensor.shape[1] + 1
+        # This is needed to don't have tall matrix have the diagonal wrap.
+        if not wrap:
+            end = tensor.shape[1] * tensor.shape[1]
+    else:
+        # 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 s[1:] != s[:-1]:
+            raise ValueError("All dimensions of input must be of equal length")
+        sz = torch.as_tensor(tensor.shape[:-1])
+        step = 1 + (torch.cumprod(sz, 0)).sum()
+
+    # Write the value out into the diagonal.
+    tensor.ravel()[:end:step] = t_val
+    return tensor
+
+
 # ### splits ###
 
 
@@ -396,6 +424,26 @@ def meshgrid(*xi_tensors, copy=True, sparse=False, indexing="xy"):
     return output
 
 
+def indices(dimensions, dtype=int, sparse=False):
+    # https://github.com/numpy/numpy/blob/v1.24.0/numpy/core/numeric.py#L1691-L1791
+    dimensions = tuple(dimensions)
+    N = len(dimensions)
+    shape = (1,) * N
+    if sparse:
+        res = tuple()
+    else:
+        res = torch.empty((N,) + dimensions, dtype=dtype)
+    for i, dim in enumerate(dimensions):
+        idx = torch.arange(dim, dtype=dtype).reshape(
+            shape[:i] + (dim,) + shape[i + 1 :]
+        )
+        if sparse:
+            res = res + (idx,)
+        else:
+            res[i] = idx
+    return res
+
+
 def bincount(x_tensor, /, weights_tensor=None, minlength=0):
     int_dtype = _dtypes_impl.default_int_dtype
     (x_tensor,) = _util.cast_dont_broadcast((x_tensor,), int_dtype, casting="safe")
@@ -620,3 +668,16 @@ def where(condition, x, y):
     else:
         result = torch.where(condition, x, y)
     return result
+
+
+# ### dot and other linalg ###
+
+def vdot(t_a, t_b, /):
+    # torch only accepts 1D arrays, numpy ravels
+    t_a, t_b = torch.atleast_1d(t_a, t_b)
+    if t_a.ndim > 1:
+        t_a = t_a.ravel()
+    if t_b.ndim > 1:
+        t_b = t_b.ravel()
+    result = torch.vdot(t_a, t_b)
+    return result

torch_np/_wrapper.py

@@ -394,6 +394,19 @@ def diag_indices_from(arr):
     return tuple(asarray(x) for x in result)
 
 
+def fill_diagonal(a, val, wrap=False):
+    tensor, t_val = _helpers.to_tensors(a, val)
+    result = _impl.fill_diagonal(tensor, t_val, wrap)
+    return asarray(result)
+
+
+@_decorators.dtype_to_torch
+def trace(a, offset=0, axis1=0, axis2=1, dtype=None, out=None):
+    tensor = asarray(a).get()
+    result = torch.diagonal(tensor, offset, dim1=axis1, dim2=axis2).sum(-1, dtype=dtype)
+    return asarray(result)
+
+
 ###### misc/unordered
 
 
@@ -483,6 +496,11 @@ def searchsorted(a, v, side="left", sorter=None):
     return arr.searchsorted(v, side=side, sorter=sorter)
 
 
+def vdot(a, b, /):
+    t_a, t_b = _helpers.to_tensors(a, b)
+    result = _impl.vdot(t_a, t_b)
+    return result.item()
+
 ###### module-level queries of object properties
 
 
@@ -615,6 +633,15 @@ def meshgrid(*xi, copy=True, sparse=False, indexing="xy"):
     return [asarray(t) for t in output]
 
 
+@_decorators.dtype_to_torch
+def indices(dimensions, dtype=int, sparse=False):
+    result = _impl.indices(dimensions, dtype=dtype, sparse=sparse)
+    if sparse:
+        return tuple(asarray(x) for x in result)
+    else:
+        return asarray(result)
+
+
 def nonzero(a):
     arr = asarray(a)
     return arr.nonzero()

torch_np/tests/numpy_tests/core/test_numeric.py

@@ -291,7 +291,6 @@ def test_take(self):
         out = np.take(a, indices)
         assert_equal(out, tgt)
 
-    @pytest.mark.xfail(reason="TODO implement trace(...)")
     def test_trace(self):
         c = [[1, 2], [3, 4], [5, 6]]
         assert_equal(np.trace(c), 5)
@@ -2627,7 +2626,7 @@ def test_exceptions(self):
         assert_raises(np.AxisError, np.rollaxis, a, 4, 0)
         assert_raises(np.AxisError, np.rollaxis, a, 0, 5)
 
-    @pytest.mark.xfail(reason="needs np.indices")
+    @pytest.mark.xfail(reason="needs fancy indexing")
     def test_results(self):
         a = np.arange(1*2*3*4).reshape(1, 2, 3, 4).copy()
         aind = np.indices(a.shape)
@@ -2833,7 +2832,6 @@ def test_outer_out_param():
     assert_equal(np.outer(arr2, arr3, out2), out2)
 
 
-@pytest.mark.xfail(reason="TODO")
 class TestIndices:
 
     def test_simple(self):

torch_np/tests/numpy_tests/lib/test_index_tricks.py

@@ -8,12 +8,12 @@
 from pytest import raises as assert_raises #, assert_raises_regex,
 
 from numpy.lib.index_tricks import (
-    mgrid, ogrid, ndenumerate, fill_diagonal,
-    ndindex, r_, s_, ix_
+    mgrid, ogrid, ndenumerate,
+    ndindex, r_, ix_
     )
 
-from torch_np import diag_indices, diag_indices_from
-from torch_np._detail._index_tricks import index_exp
+from torch_np import diag_indices, diag_indices_from, fill_diagonal
+from torch_np._detail._index_tricks import index_exp, s_
 
 
 @pytest.mark.xfail(reason='unravel_index not implemented')
@@ -358,7 +358,6 @@ def test_basic(self):
                      [((0, 0), 1), ((0, 1), 2), ((1, 0), 3), ((1, 1), 4)])
 
 
-@pytest.mark.xfail(reason='s_ not implemented')
 class TestIndexExpression:
     def test_regression_1(self):
         # ticket #1196
@@ -422,10 +421,9 @@ def test_c_():
     assert_equal(a, [[1, 2, 3, 0, 0, 4, 5, 6]])
 
 
-@pytest.mark.xfail(reason='fill_diagonal not implemented')
 class TestFillDiagonal:
     def test_basic(self):
-        a = np.zeros((3, 3), int)
+        a = np.zeros((3, 3), dtype=int)
         fill_diagonal(a, 5)
         assert_array_equal(
             a, np.array([[5, 0, 0],
@@ -434,7 +432,7 @@ def test_basic(self):
             )
 
     def test_tall_matrix(self):
-        a = np.zeros((10, 3), int)
+        a = np.zeros((10, 3), dtype=int)
         fill_diagonal(a, 5)
         assert_array_equal(
             a, np.array([[5, 0, 0],
@@ -450,7 +448,7 @@ def test_tall_matrix(self):
             )
 
     def test_tall_matrix_wrap(self):
-        a = np.zeros((10, 3), int)
+        a = np.zeros((10, 3), dtype=int)
         fill_diagonal(a, 5, True)
         assert_array_equal(
             a, np.array([[5, 0, 0],
@@ -466,7 +464,7 @@ def test_tall_matrix_wrap(self):
             )
 
     def test_wide_matrix(self):
-        a = np.zeros((3, 10), int)
+        a = np.zeros((3, 10), dtype=int)
         fill_diagonal(a, 5)
         assert_array_equal(
             a, np.array([[5, 0, 0, 0, 0, 0, 0, 0, 0, 0],
@@ -475,22 +473,22 @@ def test_wide_matrix(self):
             )
 
     def test_operate_4d_array(self):
-        a = np.zeros((3, 3, 3, 3), int)
+        a = np.zeros((3, 3, 3, 3), dtype=int)
         fill_diagonal(a, 4)
         i = np.array([0, 1, 2])
         assert_equal(np.where(a != 0), (i, i, i, i))
 
     def test_low_dim_handling(self):
         # raise error with low dimensionality
-        a = np.zeros(3, int)
-        with assert_raises_regex(ValueError, "at least 2-d"):
+        a = np.zeros(3, dtype=int)
+        with assert_raises(ValueError):
             fill_diagonal(a, 5)
 
     def test_hetero_shape_handling(self):
         # raise error with high dimensionality and
         # shape mismatch
-        a = np.zeros((3,3,7,3), int)
-        with assert_raises_regex(ValueError, "equal length"):
+        a = np.zeros((3,3,7,3), dtype=int)
+        with assert_raises(ValueError):
             fill_diagonal(a, 2)