Search code examples
pythontensorflowkerastensorflow2.0

Element wise multiplication of two list that are tf.Tensor in tensorflow


What is the fastest way to do an element-wise multiplication between a tensor and an array in Tensorflow 2?

For example, if the tensor T (of type tf.Tensor) is:

[[0, 1],  
[2, 3]]

and we have an array a (of type np.array):

[0, 1, 2]

I wand to have:

[[[0, 0],  
  [0, 0]],  
  
 [[0, 1],  
  [2, 3]],  
 
 [[0, 2],  
  [4, 6]]]  

as output.


Solution

  • This is called the outer product of two tensors. It's easy to compute by taking advantage of Tensorflow's broadcasting rules:

    import numpy as np
    import tensorflow as tf
    
    t = tf.constant([[0, 1],[2, 3]]) 
    a = np.array([0, 1, 2])
    
    # (2,2) x (3,1,1) produces the desired shape of (3,2,2)
    result = t * a.reshape((-1, 1, 1))
    # Alternatively: result = t * a[:, np.newaxis, np.newaxis]
    
    print(result)
    

    results in

    <tf.Tensor: shape=(3, 2, 2), dtype=int32, numpy=
    array([[[0, 0],
            [0, 0]],
    
           [[0, 1],
            [2, 3]],
    
           [[0, 2],
            [4, 6]]], dtype=int32)>