Efficient way of multiply 3D-tensor with matrix in Tensorflow?

Question

I have 2 3D-tensors A of shape (batch_size, dim1, dim2) and B of size (batch_size, dim2, dim3). I would like to perform the following computation:

C = np.zeros((batch_size, dim1, dim3))
for i in range(batch_size):
    c[i] = A[i].dot(B[i])

I tried C = tf.matmul(A,B), but got the error:

tensorflow.python.framework.errors_impl.InvalidArgumentError: Shape must be rank 2 but is rank 3 for 'MatMul' (op: 'MatMul') with input shapes: [10,100,100], [10,100]

Is there any efficient way to implement this in TensorFlow? Thank you!

@Roger your indicates that A and B had ranks 2 and 3. Try both with rank 3 — Maxim, Feb 22 '18 at 18:45

pfm · Answer 1 · 2018-02-22T13:29:05.683

Let's first do it with numpy:

import numpy as np
batch_size = 5
dim1 = 7
dim2 = 2
dim3 = 3

A = np.random.rand(batch_size, dim1, dim2)
B = np.random.rand(batch_size, dim2, dim3)  
C = np.zeros((batch_size, dim1, dim3)) 

for i in range(batch_size):                         
    c[i] = A[i].dot(B[i])
print(np.einsum('ikl,ilm->ikm',A,B)) - C # prints a zero array

Let's now do it using tensorflow:

import tensorflow as tf
a = tf.constant(A) # reusing numpy arrays
b = tf.constant(B)
op = tf.einsum('ikl,ilm->ikm',a,b)
with tf.Session() as sess:       
    print(sess.run(op) - C) # prints a zero array

Now according this answer, it looks like the performance should be rather similar with the naïve implementation.

Efficient way of multiply 3D-tensor with matrix in Tensorflow?

1 Answers1