Using tf.einsum with tensors of different shapes

Question

I have two tensors, A and B:

A = tf.Variable(initial_value=tf.random_normal_initializer()(shape=(1, 64, 4, 1, 1)))
B = tf.Variable(initial_value=tf.random_normal_