Stanislav Tsepa
Stanislav Tsepa

Reputation: 720

Weighted masking in TensorFlow

I have the following task: having two vectors [v_1, ..., v_n] and [w_1, ..., w_n] build new vector [v_1] * w_1 + ... + [v_n] * w_n.

For exmaple for v = [0.5, 0.1, 0.7] and w = [2, 3, 0] the result will be

[0.5, 0.5, 0.1, 0.1, 0.1].

In case of using vanilla python, the solution would be

v, w = [...], [...]
res = []
for i in range(len(v)):
    res += [v[i]] * w[i]

Is it possible to build such code within TensorFlow function? It seems to be an extension of tf.boolean_mask with additional argument like weights or repeats.

Upvotes: 0

Views: 271

Answers (1)

P-Gn
P-Gn

Reputation: 24581

Here is a simple solution using tf.sequence_mask:

import tensorflow as tf

v = tf.constant([0.5, 0.1, 0.7])
w = tf.constant([2, 3, 0])

m = tf.sequence_mask(w)
v2 = tf.tile(v[:, None], [1, tf.shape(m)[1]])
res = tf.boolean_mask(v2, m)

sess = tf.InteractiveSession()
print(res.eval())
# array([0.5, 0.5, 0.1, 0.1, 0.1], dtype=float32)

Upvotes: 1

Related Questions