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.
def generate_full_A(A): n = A.shape # 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.