How to model element wise trainable scalars for a matrix input?

Say that I want a model looks like y=A*x where X is a matrix and * is element wise product. How can I construct a tranable A?
Thank you very much

Revise:I want y=(A*B) * x where x is a vector as one input and A is a matrix as another input and B is the trainable elementwise product matrix I desire. So far, I did : = nn.Sequential(A, )
And I am confused how to step forward.

Hi @WeiHao97,

If I understand correctly you want to train your model w.r.t. B, the learning parameter matrix.

class Network(nn.Module):
   def __init__(self, shape):
       self.B = nn.Parameter(torch.zeros(shape))
       torch.nn.init.xavier_uniform_(self.B) #or any other init method
   def forward(self, x, A):
       M = A*self.B
       return x @ M.t()

shape has to be (output_dim, input_dim). input_dim being the dim of x.
Hope this is what you wanted to do.

1 Like

@spanev Thank you very much! This is exactly what I want to do.