Implement `tril_indices`, `triu_indices`, `triu_indices_from` and `tril_indices_from`

aesara/tensor/basic.py

@@ -1254,6 +1254,122 @@ def triu(m, k=0):
     )
 
 
+def tril_indices(
+    n: Union[int, ScalarVariable],
+    k: Union[int, ScalarVariable] = 0,
+    m: Optional[Union[int, ScalarVariable]] = None,
+) -> Tuple[TensorVariable, TensorVariable]:
+    """
+    Return the indices for the lower-triangle of an (n, m) array.
+
+    Parameters
+    ----------
+    n : integer scalar
+        The row dimension of the arrays for which the returned indices will be valid.
+    k : integer scalar, optional
+        Diagonal offset to use when forming the indices. `k = 0` (the default)
+        is the main diagonal, `k < 0` is below it and `k > 0` is above.
+    m : integer scalar, optional
+        The column dimension of the arrays for which the returned arrays will
+        be valid. By default m is taken equal to n.
+
+    Returns
+    -------
+    inds : tuple of TensorVariable's
+        The indices for the triangle. The returned tuple contains two arrays,
+        each with the indices along one dimension of the array.
+    """
+    return tri(n, m, k, dtype=bool).nonzero()
+
+
+def tril_indices_from(
+    a: Union[np.ndarray, TensorVariable],
+    k: Union[int, ScalarVariable] = 0,
+) -> Tuple[TensorVariable, TensorVariable]:
+    """
+    Return the indices for the lower-triangle of arr.
+
+    Parameters
+    ----------
+    arr : {array_like, TensorVariable}, shape(N, N)
+        The indices will be valid for square arrays.
+    k : integer scalar, optional
+        Diagonal offset to use when forming the indices. `k = 0` (the default)
+        is the main diagonal, `k < 0` is below it and `k > 0` is above.
+
+    Returns
+    -------
+    tril_indices_from : tuple, shape(2) of TensorVariable, shape(N)
+        Indices for the lower-triangle of arr.
+
+    Raises
+    ------
+    ValueError
+        If the input is not a 2d array.
+    """
+    if a.ndim != 2:
+        raise ValueError("The input array must be two dimensional.")
+    return tril_indices(a.shape[0], k=k, m=a.shape[1])
+
+
+def triu_indices(
+    n: Union[int, ScalarVariable],
+    k: Union[int, ScalarVariable] = 0,
+    m: Optional[Union[int, ScalarVariable]] = None,
+) -> Tuple[TensorVariable, TensorVariable]:
+    """
+    Return the indices for the upper-triangle of an (n, m) array.
+
+    Parameters
+    ----------
+    n : integer scalar
+        The row dimension of the arrays for which the returned indices will be valid.
+    k : integer scalar, optional
+        Diagonal offset to use when forming the indices. `k = 0` (the default)
+        is the main diagonal, `k < 0` is below it and `k > 0` is above.
+    m : int scalar, optional
+        The column dimension of the arrays for which the returned arrays will
+        be valid. By default m is taken equal to n.
+
+    Returns
+    -------
+    inds : tuple of TensorVariable's
+        The indices for the triangle. The returned tuple contains two arrays,
+        each with the indices along one dimension of the array.
+    """
+    return (constant(1, dtype=int) - tri(n, m, k - 1, dtype=int)).nonzero()
+
+
+def triu_indices_from(
+    a: Union[np.ndarray, TensorVariable],
+    k: Union[int, ScalarVariable] = 0,
+) -> Tuple[TensorVariable, TensorVariable]:
+    """
+    Return the indices for the upper-triangle of arr.
+
+    Parameters
+    ----------
+    arr : {array_like, TensorVariable}, shape(N, N)
+        The indices will be valid for square arrays.
+    k : integer scalar, optional
+        Diagonal offset to use when forming the indices. `k = 0` (the default)
+        is the main diagonal, `k < 0` is below it and `k > 0` is above.
+
+    Returns
+    -------
+    triu_indices_from : tuple, shape(2) of TensorVariable, shape(N)
+        Indices for the upper-triangle of arr.
+
+    Raises
+    ------
+    ValueError
+        If the input is not a 2d array.
+    """
+    if a.ndim != 2:
+        raise ValueError("The input array must be two dimensional.")
+    return triu_indices(a.shape[0], k=k, m=a.shape[1])
+
+
 class Eye(Op):
 
     __props__ = ("dtype",)
@@ -4355,4 +4471,8 @@ __all__ = [
     "full_like",
     "empty",
     "empty_like",
+    "tril_indices",
+    "tril_indices_from",
+    "triu_indices",
+    "triu_indices_from",
 ]

tests/tensor/test_basic.py

@@ -82,7 +82,11 @@ from aesara.tensor.basic import (
     tile,
     tri,
     tril,
+    tril_indices,
+    tril_indices_from,
     triu,
+    triu_indices,
+    triu_indices_from,
     unbroadcast,
     vertical_stack,
     zeros_like,
@@ -879,16 +883,34 @@ class TestTriangle:
             m_symb = matrix(dtype=m.dtype)
             k_symb = iscalar()
             f = function([m_symb, k_symb], tril(m_symb, k_symb))
+            f_indx = function(
+                [m_symb, k_symb], tril_indices(m_symb.shape[0], k_symb, m_symb.shape[1])
+            )
+            f_indx_from = function([m_symb, k_symb], tril_indices_from(m_symb, k_symb))
             result = f(m, k)
+            result_indx = f_indx(m, k)
+            result_from = f_indx_from(m, k)
             assert np.allclose(result, np.tril(m, k))
+            assert np.allclose(result_indx, np.tril_indices(m.shape[0], k, m.shape[1]))
+            assert np.allclose(result_from, np.tril_indices_from(m, k))
+            assert np.allclose(result_indx, result_from)
             assert result.dtype == np.dtype(dtype)
 
         def check_u(m, k=0):
             m_symb = matrix(dtype=m.dtype)
             k_symb = iscalar()
             f = function([m_symb, k_symb], triu(m_symb, k_symb))
+            f_indx = function(
+                [m_symb, k_symb], triu_indices(m_symb.shape[0], k_symb, m_symb.shape[1])
+            )
+            f_indx_from = function([m_symb, k_symb], triu_indices_from(m_symb, k_symb))
             result = f(m, k)
+            result_indx = f_indx(m, k)
+            result_from = f_indx_from(m, k)
             assert np.allclose(result, np.triu(m, k))
+            assert np.allclose(result_indx, np.triu_indices(m.shape[0], k, m.shape[1]))
+            assert np.allclose(result_from, np.triu_indices_from(m, k))
+            assert np.allclose(result_indx, result_from)
             assert result.dtype == np.dtype(dtype)
 
         for dtype in ALL_DTYPES:
@@ -910,6 +932,11 @@ class TestTriangle:
             check_u(m, 1)
             check_u(m, -1)
 
+        m = random_of_dtype((10,), dtype)
+        for fn in (triu_indices_from, tril_indices_from):
+            with pytest.raises(ValueError, match="must be two dimensional"):
+                fn(m)
+
 
 class TestNonzero:
     @config.change_flags(compute_test_value="raise")