Consider an nn.Linear(2,3) layer transform like the one below. It uses a 2x3 matrix of scalar weights to create a weighted sum for each scalar element in the output.

Now suppose that instead of scalars I want to operate on 1d vectors (e.g. word embeddings). So in this case a and b would be vectors and the weights w would each be square matrices.

Can one of the pre-built layer types in nn accomplish this (for linear or convolutions)? Or is there some way to transform my data to accomplish the equivalent of this with the standard layer types?


EDIT: I'm aware that for a given layer like the example I could just do all of the permutations of matrix operations and store the weights for them myself. What I'm wondering is if there is a better way.

enter image description

EDIT: Ok, I think what I want does not exist but let me try to clarify on the hopes that maybe there is some similar operation...

Ignoring the bias vector nn.Linear(2,3) is doing the following:

What I'm looking for is the analog of this where a, b, c, ... are vectors and the weighting is a linear transformation of those vectors. I apologize if my notation is off here but, written as row vectors:

enter image description here

This means that we'd need a 3 dimensional set of weights, not a 2d as in nn.Linear. Again, I guess this is just a different operation entirely but I'm wondering if this kind of thing has a name or if there is some reason it isn't normally needed (e.g. some transformation of the data that really is equivalent to this).


1 Answer 1


If I understand your question1, you'd like an input tensor of shape (2, 3), where each element is actually a vector. If that is the case, you are essentially talking about a 3d tensor (instead of 2d).

For example, if each vector is 5 elements long, your input would have the shape (2, 3, 5).

To compute convolutions over this, there are layers that take the dimensions as parameters - have a look at the Convolutional layers like Conv3d.

To compute a linear layer on this input, you still just need to flatten or reshape the tensor to be a single vector.

If you vector lengths can vary, you will have to consider making them all match by either truncating them (making some shorter), or padding (adding e.g. zeros to make shorter ones longer) - or a mixture of both techniques.

1 If I haven't fully understood, please add more details to the example in your question


Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Not the answer you're looking for? Browse other questions tagged or ask your own question.