rosefun
rosefun

Reputation: 1857

How to use tf.gather if the index is probabilistic form in Tensorflow?

I want to use probability to choose the array.

import tensorflow as tf
from keras import backend as K
T = np.array([[0.6, 0.1, 0.3], [0.2,0.7,0.1], [0.1,0.1,0.8]])
labels = np.array([1,2,0])
preds = np.array([[0.2,0.7,0.1]])

In tf.gather, if index is integer, it can work.

K.get_value(tf.gather(T, labels))

array([[0.2, 0.7, 0.1],
       [0.1, 0.1, 0.8],
       [0.6, 0.1, 0.3]])

However, now my index is probabilistic form, I hope it can still work like this:

K.get_value(tf.gather(T, preds))

Here, I want to select the index 0 with a probability of 0.2, select the index 1 with a probability of 0.7, and select the index 2 with a probability of 0.1.

Expected result:

0.2*[0.6, 0.1, 0.3] + 0.7*[0.2,0.7,0.1] + 0.1*[0.1,0.1,0.8]

Upvotes: 1

Views: 52

Answers (1)

thushv89
thushv89

Reputation: 11333

What you're doing is a type of weighted sum so the good news is that you don't need to use tf.gather. In general, you need the following.

Given T (n, t_cols) and preds (p_rows, n), you need, (p_rows, t_cols) sized output. Where each row in the output is a weighted sum of rows of T.

Here's what you can do. This involves understanding numpy broadcasting rules.

We are doing the following.

  • Make T of shape [t_cols, n, 1]
  • Make preds of shape [n, p_rows]
  • Multiply the two, which gives [t_cols, n, p_rows]
  • Reduce sum on n, which gives [t_cols, p_rows]
  • Transpose, which gives [p_rows, t_cols]
T = tf.constant(np.array([[0.6, 0.1, 0.3], [0.2,0.7,0.1], [0.1,0.1,0.8]]))
preds = tf.constant(np.array([[0.2,0.7,0.1],[0.2,0.4,0.4]]))

mul = tf.transpose(T)[:,:,tf.newaxis] * tf.transpose(preds)
res = tf.transpose(tf.reduce_sum(mul, axis=1))

res will have the result you need.

Upvotes: 1

Related Questions