CODEWITHSUNDEEP
pythontensorflowtensor
Jsevillamol
Jsevillamol

Reputation: 2563

Tensor contraction in tensorflow

I have a tensor weights of shape (?,4) and a tensor embeddings of shape (?,4,1024).

I would like to contract the tensor by taking a weighted mean of the 4 tensors in each row of embeddings according to the corresponding weights, finally producing a tensor output of shape (?,1024).

How can I do that? I tried with output = tf.tensordot(weights, embeddings, axes = [1,1]) but that produced a tensor of shape (?,?,1024) instead.

Upvotes: 1

Views: 535

Answers (1)

javidcf
javidcf

Reputation: 59731

You can do that like this:

import tensorflow as tf

weights = tf.placeholder(tf.float32, [None, 4])
embeddings = tf.placeholder(tf.float32, [None, 4, 1024])
output = tf.einsum('ij,ijk->ik', weights, embeddings)

You can express the same thing through matrix product, not sure if there would be any difference in performance:

output = tf.squeeze(tf.expand_dims(weights, 1) @ embeddings, 1)

You could also just multiply and reduce, although that would in principle have worse performance due to having an intermediate tensor.

output = tf.reduce_sum(tf.expand_dims(weights, 2) * embeddings, axis=1)

Upvotes: 3

Related Questions