import torch


def Ux(diags, offdiags, x):
    """
    Let U be an upper block-bidiagonal matrix whose
    - diagonals are given by diags
    - upper off-diagonals are given by offdiags
    We would like to compute U@x
    """
    n = diags.shape[0]
    m = offdiags.shape[0]

    # non-square matrix
    if n == m:
        return torch.einsum("ijk,ik...->ij", diags, x[:-1]) + torch.einsum(
            "ijk,ik...->ij", offdiags, x[1:]
        )
    else:
        leaf1 = torch.einsum("ijk,ik...->ij", diags, x)
        leaf2 = torch.einsum("ijk,ik...->ij", offdiags, x[1:])
        leaf_sum = leaf1[:-1] + leaf2
        # remembering the final row of the matrix which just has a diagional element and not an upper diagional element
        return torch.cat([leaf_sum, leaf1[-1].unsqueeze(0)], dim=0)