A fully connect layer which maps a vector to a matrix

My network input has dimension $[K, n]$ where $K$ is the number of samples, I want to create a layer that will map this vector of dimension $n$ to a matrix $[d, n]$, that is, the output of the network is $[K, d, n]$ does anyone know which nn.layer I can use? Thank you very much!

Hi Zecheng!

You can simply take the output of a Linear that has out_features = d * n
and view() (or reshape() or Unflatten) it as a batch of d x n matrices:

>>> import torch
>>> torch.__version__
'1.11.0'
>>> K = 2
>>> n = 5
>>> d = 7
>>> input = torch.randn (K, n)
>>> torch.nn.Linear (in_features = n, out_features = d * n) (input).view (-1, d, n).shape
torch.Size([2, 7, 5])

Because Linear is fully connected, it has the right number degrees of
freedom to generate any affine (i.e., “linear”) transformation that maps
a vector of length n to a matrix of size d x n. (This transformation then
gets applied sample-wise to a batch of K samples.)

Best.

K. Frank