Multilinear map

[santient]

🚀 Feature

Multilinear layer: a generalization of the Linear and Bilinear layers/functions to any configurable number of variables. This should be a learnable multilinear map that does the equivalent of this function, plus a bias (from the Multilinear map Wikipedia page):

Details

Implement the following:

torch.nn.functional.multilinear((input1, ..., inputN), weight, bias=None)
torch.nn.Multilinear((in1_features, ..., inN_features), out_features, bias=True)

I think the rank of the weight tensor would be the number of input features plus one.

Motivation

This is useful for situations where you want to combine more than two vectors using a learnable function that is linear in all variables. A good use case would be tensor fusion with more than two variables.

I plan to implement this for use in a project. It should be fairly easy to implement using something like torch.einsum, but a native C++ implementation could speed things up significantly.

Should I submit a PR when this is done?

cc @albanD @mruberry

[ezyang]

FWIW, einsum is implemented in C++. (Tensor networks may be relevant). Also cc @zou3519

[colesbury]

@ezyang and @fmassa suggest implementing this with einsum for now.

[acturner]

@santient Is this implemented? If so, could you post a link to your implementation or PR?

[santient]

@acturner Forgot about this issue but I just wrote up a quick implementation here. I found no need to use einsum. Lmk if you notice a mistake or any bugs!

[acturner]

@acturner Forgot about this issue but I just wrote up a quick implementation here. I found no need to use einsum. Lmk if you notice a mistake or any bugs!

Thanks so much!

[santient]

Just FYI, I updated my implementation using einsum to support batched vectors