Not sure if this is a feature but is it possible to make individual linear parameter for each input along the axis that is not the input channel? say an Nx4 matrix fed to an nn.linear layer of size 4 input and 4 for output, I would have a parameter size of Nx4 for bias and Nx4x4 weight instead of 4 and 4x4. (I know I can do this manually by setting up my own parameters but is there a way to do so in nn.linear?)
just do it the way you think how it might work