Keras - bilinear data transformation (like in pytorch)
So I'm trying to implement in Keras (tf) &quot;bilinear transformation to the incoming data&quot; (coming from pytorch). Transformation is defined in torch.nn.Bilinear as
y = x_1^T A x_2 + b 
I built custom layer in Keras that in method call() will return transformed data (I'd like to use this layer as a part of my model later on). I do have however problem with implementation of the transformation function itself due to 3d input I have.
and shapes are as follows:

x_1= TensorShape([72, 10, 6])
x_2= TensorShape([72, 10, 6])
self.w = TensorShape([24, 6, 6]) (24 coming from defined out_features=4*feature_vector)
self.b = TensorShape([24])

output I'd like to use as an input for LSTM layer.
So my formula in tf looks like that:
a = tf.matmul(tf.matmul(input1, self.w), input2) + self.b
and I tried transposing different parts but it doesn't work - I have a problem with not compatible dimensions and I'd really appreciate any hint

I have a detailed implementation of bilinear with pytorch and you can translate to keras easily
a = torch.randn(2,3) # input1
b = torch.randn(3,7,5) # bilinear matrix weight (input1 dim, input2 dim, output dim)
c = torch.randn(2,7) # input2

B = nn.Parameter(b)
t = a @ B.permute(2,0,1)
output = (t * c).sum(-1).t()

