Understanding batch_dot() in Keras with Tensorflow backend

Does the first dimension (with value ?) get counted for the arguments of axes?

Yes, it is counted.

The fact is that the first dimension in Input layer, in the above example, is batch size while in K.ones() it is not. As a result, axes [3, 3] for Input is equal to axes [2, 2] in K.ones(). In your code, the following two batch_dot are equal:

q = tf.keras.layers.Input(shape=(36, 8, 24)) # shape =(?, 36,8,24)
k = tf.keras.layers.Input(shape=(36, 8, 24)) # shape =(?, 36,8,24)
print(tf.keras.backend.batch_dot(q, k, axes=[3, 3]))

q = tf.keras.backend.ones(shape=(36, 8, 24))
k = tf.keras.backend.ones(shape=(36, 8, 24))
print(tf.keras.backend.batch_dot(q, k, axes=[2, 2]))

Note that, in K.ones(), if the shape was symbolic, we cannot return a variable, and will return a dynamically-shaped tensor instead. What does it mean? See following example for better understanding:

a = tf.keras.layers.Input(shape=(30,))
c = tf.keras.backend.ones(shape=tf.shape(a))
print(c) # shape=(?, 30)
d = tf.keras.backend.ones(shape=(30, 40))
print(d) # shape=(30,40)

Output dimensions from batch_dot() operation is unexpected

K.batch_dot(q,k,axes=[1,1])
<tf.Tensor 'MatMul_11:0' shape=(?, 36, 24, 24) dtype=float32>
K.batch_dot(q,k,axes=[2,2])
<tf.Tensor 'MatMul_12:0' shape=(?, 36, 24, 24) dtype=float32>

Why on earth this happens while axes are different?

To answer this question, we should know about the underlying implementation of batch_dot. If the rank of input tensors is not 2, then our batch_dot behaves as tf.matmul operation that one of the input tensors is conjugate transposed. As a result, when our input tensors have a rank of 3 and we set the axis, to 0 or 1 they calculated the same things but when set axes to 2 it calculated something different:

a = np.array([[[1, 2, 3],
               [3, 2, 1]]])  # rank 3

b = np.array([[[1, 3, 3],
               [2, 2, 0]]])  # rank 3

a = tf.constant(a, dtype=tf.float32)
b = tf.constant(b, dtype=tf.float32)

c = tf.matmul(a, b, adjoint_a=True, adjoint_b=False)  # when axes is [0,0] or [1,1]
d = tf.matmul(a, b, adjoint_a=False, adjoint_b=True)  # when axes is [2,2]
print(c.shape)  # shape=(1,3,3)
print(d.shape)  # shape=(1,2,2)

The same things happened in your example:

a = np.array([[[1, 2, 3],
               [3, 2, 1]]])

b = np.array([[[1, 3, 3],
               [2, 2, 0]]])

q = tf.placeholder(dtype=tf.float32, shape=(None, 2, 3))  
k = tf.placeholder(dtype=tf.float32, shape=(None, 2, 3))  
res1 = tf.keras.backend.batch_dot(q, k, axes=0)
res2 = tf.keras.backend.batch_dot(q, k, axes=1)
res3 = tf.keras.backend.batch_dot(q, k, axes=2)
with tf.Session() as sess:
    feed_dic = {q: a, k: b}
    print(sess.run(res1, feed_dict=feed_dic))
    print(20 * '-')
    print(sess.run(res2, feed_dict=feed_dic))
    print(20 * '-')
    print(sess.run(res3, feed_dict=feed_dic))

Things will be clear if you look at the source code at https://github.com/tensorflow/tensorflow/blob/a6d8ffae097d0132989ae4688d224121ec6d8f35/tensorflow/python/keras/backend.py#L1437

we can directly go line1507

if ndim(x) == 2 and ndim(y) == 2:
    if axes[0] == axes[1]:
      out = math_ops.reduce_sum(math_ops.multiply(x, y), axes[0])
    else:
      out = math_ops.reduce_sum(
          math_ops.multiply(array_ops.transpose(x, [1, 0]), y), axes[1])
else: 
    adj_x = None if axes[0] == ndim(x) - 1 else True
    adj_y = True if axes[1] == ndim(y) - 1 else None
    out = math_ops.matmul(x, y, adjoint_a=adj_x, adjoint_b=adj_y)

As it shows, it only checks the adj_x and adj_y and doesn't pass the axes parameter to math_ops.matmul method. That is the reason you are getting the same result when axes is [1,1] and [2,2].

We can use below code to verify:

q = K.ones(shape=range(1, 10))
k = K.ones(shape=range(1, 10))
for i in range(10): print(i, K.batch_dot(q,k,axes=[i,i]))

it will print

0 Tensor("MatMul_7:0", shape=(1, 2, 3, 4, 5, 6, 7, 9, 9), dtype=float32)
1 Tensor("MatMul_8:0", shape=(1, 2, 3, 4, 5, 6, 7, 9, 9), dtype=float32)
2 Tensor("MatMul_9:0", shape=(1, 2, 3, 4, 5, 6, 7, 9, 9), dtype=float32)
3 Tensor("MatMul_10:0", shape=(1, 2, 3, 4, 5, 6, 7, 9, 9), dtype=float32)
4 Tensor("MatMul_11:0", shape=(1, 2, 3, 4, 5, 6, 7, 9, 9), dtype=float32)
5 Tensor("MatMul_12:0", shape=(1, 2, 3, 4, 5, 6, 7, 9, 9), dtype=float32)
6 Tensor("MatMul_13:0", shape=(1, 2, 3, 4, 5, 6, 7, 9, 9), dtype=float32)
7 Tensor("MatMul_14:0", shape=(1, 2, 3, 4, 5, 6, 7, 9, 9), dtype=float32)
8 Tensor("MatMul_15:0", shape=(1, 2, 3, 4, 5, 6, 7, 8, 8), dtype=float32)
9 Tensor("MatMul_16:0", shape=(1, 2, 3, 4, 5, 6, 7, 9, 9), dtype=float32)

Except when i is 8, all others are returning the same result.