Lower-triangular circulant embedding of pytorch vector

lowtr_circulant.py

"""
Partially inspired from https://stackoverflow.com/a/70686229
"""
import torch as pt

def lowtr_circulant(diag: pt.Tensor, dim: int=-1):
    r"""
    From a vector ``diag``, computes its flipped circular embedding with zeros
    above the diagonal.
    Can be batched in left dimension(s).
    
    ::math
        \lambda \to \begin{matrix}
            \lambda_1 &        &           & \\
            \lambda_2 & \ddots & 0         & \\
            \vdots    & \cdot  & \ddots    & \\
            \lambda_n & \dots  & \lambda_2 & \lambda_1 \\
        \end{matrix}
    """
    d = diag.shape[dim] # Dimension of the vector
    zero_padds = pt.zeros(d - 1).to(diag)
    return pt.cat([zero_padds, diag]).unfold(dim, d, 1).flip((-1,))