How to assemble layers to perform different linear transformations on each column of matrix?

Question

I have a matrix $A_{m×n}$ and I want to assemble a network that performs $n$ different trainable linear transformations $T_{p×m}(i), i\in {1,2,...n}$ on its corresponding columns, (i.e. $B^i = T(i)A^i)$. However, the LinearLayer[] in mathematica will apply to all the columns. It's possible to use n PartLayer[] to extract every column in $A$, but that may not work well when $n$ is big. Is there any easy way to implement this?
My solution looks like this: First time using StackExchange, thank you so much :)

Answer 1

This might not be the most compact or elegant to represent the NetGraph, but at least the size of the diagram is constant no matter how large the parameters n,m and p are.

trynet = With[{p = 5, n = 200, m = 100}, With[{},
   NetGraph[<|
     "WeightsLayer" -> 
      NetArrayLayer["Output" -> {Automatic, p, Automatic}],
     "Transpose" -> TransposeLayer[],
     "T Layer" -> 
      NetMapThreadOperator[
       NetGraph[{NetMapThreadOperator[
          ThreadingLayer[Times, 
           InputPorts -> <|"Weights" -> Automatic, 
             "Input" -> Automatic|>], <|"Weights" -> 1|>], 
         AggregationLayer[Total, 2]}, {NetPort["Input"] -> 
          NetPort[1, "Input"], 
         NetPort["Weights"] -> NetPort[1, "Weights"], 
         1 -> 2}], {"Input" -> 1, "Weights" -> 1}]
     |>, {
     "WeightsLayer" -> NetPort["T Layer", "Weights"],
     NetPort["Input"] -> "Transpose" -> NetPort["T Layer", "Input"]
     }, "Input" -> {m, n}]]]

Answer 2

Code autogeneration ^_^

dims = {10, 5};

n = 8; (*LinearLayer*)

NetGraph[
 Join[Table[{PartLayer[{All, i}], LinearLayer[n]}, {i, dims[[2]]}], {CatenateLayer[]}],
 Join[Table[NetPort["Input"] -> i, {i, dims[[2]]}], {Range[dims[[2]]] -> dims[[2]] + 1}],
 "Input" -> dims
 ]