How to create Sparse Variables with Variable values?

1

Hi!
I have some matrix (n, 3) Variable A, which defines a tridiagonal matrix.
After some computations I want to create a dense (n, n) matrix.
Something like

def generate_full_A(A):
    n = A.shape[0]
    # TODO: rewrite it
    full_A = torch.diag(A[:, 1])
    for i in range(n - 1):
        full_A[i, i + 1] = A[i, 2]
    for i in range(1, n):
        full_A[i, i - 1] = A[i, 0]
    return full_A

works, but is too slow. Is there any way to perform this operation without loosing “Variable structure”? I want to compute gradients of the whole program later.

I’ve tried torch.sparse, but as far as I can see it is not able to work with Variable values.