[FIXED] How to multiply Tensorflow arrays across specified indicies


I would like to multiply two Tensorflow Arrays in a certain way as shown in the code below:

import tensorflow as tf

from tensorflow.keras import mixed_precision
policy = mixed_precision.Policy('mixed_float16')
print('Compute dtype: %s' % policy.compute_dtype)
print('Variable dtype: %s' % policy.variable_dtype)

a = tf.random.normal(shape=[1000, 1439])
b = tf.random.normal(shape=[1000, 1439])

for i in range(0,b.shape[0]):

Is there a more elegant and shorter way of doing this kind of multiplication without loops? Also I would like to have the final product in a single Tensorflow array rather than in a list. In Numpy, I can achieve the above with the following commands but somehow it doesnt work with Tensorflow arrays:

np.einsum("ij, kj->ikj", a, b)


a.reshape(a.shape[0],1,a.shape[1]) * ([b]*a.shape[0])



tf.unstack(tf.einsum("ij, kj->ikj", a, b))

should give you the same result as your example when using the same a and b. Without unstack, you will have a tensor with the shape (1000, 1000, 1439), but note that these tensors will take up a lot of memory and your program will probably crash.

Answered By – AloneTogether

Answer Checked By – Senaida (Easybugfix Volunteer)

Leave a Reply

(*) Required, Your email will not be published