How to multiply Tensorflow arrays across specified indicies-CodePudding

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')
mixed_precision.set_global_policy(policy)
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])

final_product=[]
for i in range(0,b.shape[0]):
    product=a[i,:]*b
    final_product.append(product)

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)

or

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

CodePudding user response：

Running:

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.