Tensor multiply along axis in tensorflow

Question

The problem I have using tensorflow is as follows:

For two tensors

[[x11,x12...],[x21,x22...],...[xn1,xn2...]]

and

[y1,y2,...yn],

I want to multiply them along axis 0 to get

[[x11*y1,x12*y1...],[x21*y2,x22*y2...]...]

For example, for [[1,2],[3,4]] and [1,2], I want to get the result tensor [[1,2],[6,8]].

The real scenario is that I have two tensors A and B shaped (batches,height,width,n_channels) and (batches,1). Both are tensors defined in tensorflow. For each image of A in the batch I want to multiply it with the corresponding value in B.

Answer 1

Given a 2-dimensional tensor x and a vector y, you just need to do:

result = x * tf.expand_dims(y, axis=-1)

Or, if you like it more:

result = x * y[:, tf.newaxis]