Efficient way to set tensor values using lower triangular indices


I’m looking for an efficient way to properly set batched (size = 2) lower triangular tensors each with 3 x 3 dimensions. Lets call this tensor Sigma [2, 3, 3]. I’ve created parameters Theta [2, 6] which will be used to fill Sigma. Note that 6 parameters are required to fill/set a lower triangular matrix with dim = 3 (e.g. (3*(3+1))/2). Currently my code is:

Theta = torch.nn.Parameter(torch.Tensor(2, 6))
torch.nn.init.xavier_uniform_(Theta, gain=1.414)
Sigma = torch.nn.Parameter(torch.Tensor(2, 3, 3))
tril_indices = torch.tril_indices(row=3, col=3, offset=0)
Sigma[0, tril_indices[0], tril_indices[1]] = Theta[0]
Sigma[1, tril_indices[0], tril_indices[1]] = Theta[1]

The above code works but is there a more efficient way that I replace the last two lines with a single line of code? If my batch size is say 16 I don’t want to have 16 separate lines of code setting the appropriate tensor values.

Ideally, I’d like Sigma[index_tensor] = Theta. But because Theta and Sigma have different shapes I am getting a errors for broadcasting. I’ve also tried thinking about using scatter_ but the function seems like it can’t be used for my use case.


Did you find a solution for it?

You could always use a loop to set it but I was wondering if someone found a better solution.