Python pair-wise multiplication using broadcasting on the first dimension

Question

I have one-hot vector of shape 1 * n

v= [0.0, 1.0, 0.0] for n = 3

and a matrix of shape n * m * r as( m and r can be any number but the first dimention is as n)

m = [[[1,2,3,],[4,5,6]], [[5,6,7],[7,8,9]], [[2,4,7],[1,8,9]]]

I want to multiple a * b using a broadcasting mechanism such that only the sub-matrix corresponding to 1.0 element in vector v is kept in the multiplication of v * m and all other sub-matrices become zero ( because all other elements are zero in v) as:

prod = [[[0,0,0],[0,0,0]], [[5,6,7],[7,8,9]] , [[0,0,0],[0,0,0]]]

Answer 1

In Tensorflow, you basically want to add the extra dimensions to the end of dimensions, otherwise the broadcast happens in the last dimension. So this code works like you want:

import tensorflow as tf

sess = tf.Session()

v= [0.0, 1.0, 0.0]
vT = tf.expand_dims( tf.expand_dims( tf.constant( v, dtype = tf.float32 ), axis = -1 ), axis = -1 )
m = [[[1,2,3],[4,5,6]], [[5,6,7],[7,8,9]], [[2,4,7],[1,8,9]]]
mT = tf.constant( m, dtype = tf.float32 )
prodT = vT * mT
#prod = [[[0,0,0],[0,0,0]], [[5,6,7],[7,8,9]] , [[0,0,0],[0,0,0]]]

sess.run( prodT )

Output:

array([[[0., 0., 0.], [0., 0., 0.]], [[5., 6., 7.], [7., 8., 9.]], [[0., 0., 0.], [0., 0., 0.]]], dtype=float32)