Trainable vector and matrix in the same layer

Hello,

I have the following linear layer that l want to add to my model such that w_{k} and A are respectively learnable vector and matrix. With X[50,18] and A[50,50]

eq

How can l learn w_{k} with different k and A in the same layer ?

I get stuck at writing in pytorch forward, backward.

Thank you for your help