Hi, I’m wondering how can I do an outer product of two tensor, of shapes (batch_size, dim), (batch_size, dim), so that the outer product is applied in the last dimension, resulting in a tensor of shape (batch_size, dim, dim), according to the formula:

x_i_j times y_i_k => a_i_j_k = x_i_j y_i_k

Don’t know how to use latex here, I hope that’s clear. If the batch size was one, there would be two vectors of which I want the outer product. Since there is a batch, I would need to do the same for every sample of the batch.

In general, how can operations that don’t broadcast (like outer) be applied to specific dimensions of higher dimensional tensors, especially in the case you have an extra dimension (like the batch size) that won’t be involved in the operation?

Answering my own question, einsum seems to fit perfectly with the problem (I can even use the same schematic notation I used in the question lol). If there is some way that is more specific then someone might feel free to comment.