user1305241
user1305241

Reputation: 53

Indexing tensor with another tensor

I have a tensor xi of shape (?, 20, 10) and another tensor y_data of shape (?, 20, 1). I want to use the y_data tensor to "index" the xi tensor in order to do something like tf.exp(xi[y_data] - tf.log(tf.reduce_sum(xi, axis=2)).

E.g. tf.exp(xi[:, :, 4] - tf.log(tf.reduce_sum(xi, axis=2))) results in a tensor of shape (?, 20). I just want to get the index, here 4, out ot another tensor.

Thanks in advance!

Upvotes: 1

Views: 245

Answers (1)

Sunreef
Sunreef

Reputation: 4542

In this case, I would use a loop on the possible values for y_data which I will assume go from 0 to 9.

result = tf.zeros(tf.shape(y_data), tf.float32)
for i in range(10):
    result = tf.where(tf.equal(y_data, i), tf.exp(xi[:, :, i:i+1]), result)
result = tf.reshape(result, [-1, 20])
result -= tf.log(tf.reduce_sum(xi, axis=2))

Probably not the most efficient but that's the only way I could think of.

Upvotes: 1

Related Questions