I would like to do a matrix multipication which I am not sure of how to implement. It would be an implementation of a doing a different linear forward for every 2D element in the batch.
For every 2D element of shape [seq_len, hidden_in] in the batch I would like to multiply with a specific matrix of shape [hidden_out, hidden_in] to create a batch output of elements of shape [seq_len, hidden_out]. So, the term batch here is a bit misleading as usually you would like to perform the same operation on all elements in the batch, but I would like to do a different operation for each element in the batch, but the same operation along the sequence dimension.
The (batched) shapes are as follows:
[batch, seq_len, hidden_in]
[batch, hidden_out, hidden_in]
Desired output shape:
[batch, seq_len, hidden_out]
It would replace an implementation with linear layers in a loop, something along the lines of:
# weights_for_elements_in_batch is of shape [batch, hidden_out, hidden_in] # input is of shape [batch, seq_len, hidden_in] outs =  for i in range(batch_size): out = torch.nn.functional.linear(input[i, :, :], weight=weights_for_elements_in_batch[i, :, :]) output = torch.stack(outs, dim=0) # stack all elements in batch
If anyone has a suggestion, please let me know. Help is appreciated!