The first operation will use a parameter of shape [1, 1] and an input of [2, 3]:
self.w * x

So it seems that the feature dimension isn’t even matching.
Could you explain your use case a bit?
self.W has the shape [output_dim, output_dim], which also doesn’t consider the input feature dimension (but isn’t used at all).