batch_dot with variable batch size in Keras

Question

I'm trying to writting a layer to merge 2 tensors with such a formula

The shapes of x[0] and x[1] are both (?, 1, 500).

M is a 500*500 Matrix.

I want the output to be (?, 500, 500) which is theoretically feasible in my opinion. The layer will output (1,500,500) for every pair of inputs, as (1, 1, 500) and (1, 1, 500). As the batch_size is variable, or dynamic, the output must be (?, 500, 500).

However, I know little about axes and I have tried all the combinations of axes but it doesn't make sense.

I try with numpy.tensordot and keras.backend.batch_dot(TensorFlow). If the batch_size is fixed, taking a = (100,1,500) for example, batch_dot(a,M,(2,0)), the output can be (100,1,500).

Newbie for Keras, sorry for such a stupid question but I have spent 2 days to figure out and it drove me crazy :(

    def call(self,x):
            input1 = x[0]
            input2 = x[1]
            #self.M is defined in build function
            output = K.batch_dot(...)
            return output

Update:

Sorry for being late. I try Daniel's answer with TensorFlow as Keras's backend and it still raises a ValueError for unequal dimensions.

I try the same code with Theano as backend and now it works.

>>> import numpy as np
>>> import keras.backend as K
Using Theano backend.
>>> from keras.layers import Input
>>> x1 = Input(shape=[1,500,])
>>> M = K.variable(np.ones([1,500,500]))
>>> firstMul = K.batch_dot(x1, M, axes=[1,2])

I don't know how to print tensors' shape in theano. It's definitely harder than tensorflow for me... However it works.

For that I scan 2 versions of codes for Tensorflow and Theano. Following are differences.

In this case, x = (?, 1, 500), y = (1, 500, 500), axes = [1, 2]

In tensorflow_backend:

return tf.matmul(x, y, adjoint_a=True, adjoint_b=True)

In theano_backend:

return T.batched_tensordot(x, y, axes=axes)

(If following changes of out._keras_shape don't make influence on out's value.)

Answer 1

Your multiplications should select which axes it uses in the batch dot function.

• Axis 0 - the batch dimension, it's your ?
• Axis 1 - the dimension you say has length 1
• Axis 2 - the last dimension, of size 500

You won't change the batch dimension, so you will use batch_dot always with axes=[1,2]

But for that to work, you must ajust M to be (?, 500, 500).
For that define M not as (500,500), but as (1,500,500) instead, and repeat it in the first axis for the batch size:

import keras.backend as K

#Being M with shape (1,500,500), we repeat it.   
BatchM = K.repeat_elements(x=M,rep=batch_size,axis=0)
#Not sure if repeating is really necessary, leaving M as (1,500,500) gives the same output shape at the end, but I haven't checked actual numbers for correctness, I believe it's totally ok. 

#Now we can use batch dot properly:
firstMul = K.batch_dot(x[0], BatchM, axes=[1,2]) #will result in (?,500,500)

#we also need to transpose x[1]:
x1T = K.permute_dimensions(x[1],(0,2,1))

#and the second multiplication:
result = K.batch_dot(firstMul, x1T, axes=[1,2])