CODEWITHSUNDEEP
pythontensorflow
QuadraticOne
QuadraticOne

Reputation: 57

Tensorflow: zip over first dimension of two tensors of different shape

Let's say I have two tensors, whose shapes are [b, n] and [b, n, m] respectively. These can be interpreted as a batch of input vectors each of shape [n] and a batch of weight matrices each of shape [n, m], where the batch size is b. I would like to pair these up element-wise across the first dimension, so each input vector has a corresponding weight matrix, and then multiply each input by its weights, resulting in a tensor of shape [b, m].

In normal Python I suspect this would look something like

output_list = [matmul(w, i) for w, i in zip(weight_list, input_list)]

but haven't been able to find a Tensorflow analogue; is there a way of doing this?

Upvotes: 1

Views: 642

Answers (1)

giser_yugang
giser_yugang

Reputation: 6166

tf.matmul can do a matmul over each training example in the batch. But you need to deal with some dimensions problem to achieve your goal.

import tensorflow as tf

b,n,m = 4,3,2
weight_list = tf.random.normal(shape=(b,n,m))
input_list = tf.random.normal(shape=(b,n))
result = tf.squeeze(tf.matmul(tf.expand_dims(input_list,axis=1),weight_list))
print(result.shape)

(4, 2)

Upvotes: 1

Related Questions