Broadcast tf.matmul with dynamic shapes

Question

I would like to broadcast a tf.matmul operation between two tensors of ranks 2 and 3, one of which contains an "unknown" shaped dimension (basically a 'None' value in a particular dimension).

The problem is that with dynamic dimensions tf.reshape and tf.broadcast_to don't seem to work.

x = tf.placeholder(shape=[None, 5, 10], dtype=tf.float32)
w = tf.ones([10, 20])
y = x @ w
with tf.Session() as sess:
  r1 = sess.run(y, feed_dict={x: np.ones([3, 5, 10])})
  r2 = sess.run(y, feed_dict={x: np.ones([7, 5, 10])})

Take the above code as an example. In this case I'm feeding two different batches of 3 and 7 elements each. I would like r1 and r2 to be the result of matrix-multiplying w by each of the 3 or 7 elements from these batches. Therefore the resulting shapes for r1 and r2 respectively would be (3, 5, 20) and (7, 5, 20), but instead I'm getting:

ValueError: Shape must be rank 2 but is rank 3 for 'matmul' (op: 'MatMul') with input shapes: [?,5,10], [10,20].

Answer 1

w can be expanded to a rank-3 tensor with a batch-size equal to that of the input. Then, the matmul operation can be performed

x = tf.placeholder(shape=[None, 5, 10], dtype=tf.float32)
w = tf.ones([10, 20])

number_batches = tf.shape(x)[0]
w = tf.tile(tf.expand_dims(w, 0), [number_batches, 1, 1])
y = x @ w
with tf.Session() as sess:
  print(sess.run(y, feed_dict={x: np.ones([2, 5, 10])}))
  print(sess.run(y, feed_dict={x: np.ones([3, 5, 10])}))

live code here