Application of nn.Linear layer in pytorch on additional dimentions
How is the fully-connected layer (nn.Linear) in pytorch applied on "additional dimensions"? The documentation says, that it can be applied to connect a tensor (N,*,in_features) to (N,*,out_features), where N in the number of examples in a batch, so it is irrelevant, and * are those "additional" dimensions. Does it mean that a single layer is trained using all possible slices in the additional dimensions or are separate layers trained for each slice or something yet different?

There are in_features * out_features parameters learned in linear.weight and out_features parameters learned in linear.bias. You can think of nn.Linear working as


reshape the tensor to some (N', in_features), where N' is the product of N and all dimensions described with *: input_2d = input.reshape(-1, in_features)
Apply a standard matrix-matrix multiplication output_2d = linear.weight @ input_2d.
Add the bias output_2d += linear.bias.reshape(1, in_features) (notice we broadcast it across all N' dimensions)
Reshape the output to have the same dimensions as input, aside from the last one: output = output_2d.reshape(*input.shape[:-1], out_features)
return output


So the leading dimension N is treated the same as the * dimensions. The documentation makes N explicit to let you know that the input has to be at least 2d, but can be as many dimensional as you wish.
