Getting the indices of several elements in a Tensorflow at once

user_6657314 0
Viewed 126

i'm new in Tensorflow.

i have one question.

there is 1d array here.

 values = [101,103,105,109,107]

 target_values = [105, 103]

I want to get an indices about target_values from values at once.

Indices extracted from the example above will be shown below. 

indices = [2, 1]

when i using tf.map_fn function. This problem can be solved easily.

# if you do not change data type from int64 to int32. TypeError will riase
values = tf.cast(tf.constant([100, 101, 102, 103, 104]), tf.int64)
target_values = tf.cast(tf.constant([100, 101]), tf.int64)
indices = tf.map_fn(lambda x: tf.where(tf.equal(values, x)), target_values)

thank you!

tensorflow python
1 Answers

Assuming all values in target_values are in values, this is one simple way to do that (TF 2.x, but the function should work the same for 1.x):

import tensorflow as tf

values = [101, 103, 105, 109, 107]
target_values = [105, 103]

# Assumes all values in target_values are in values
def find_in_array(values, target_values):
    values = tf.convert_to_tensor(values)
    target_values = tf.convert_to_tensor(target_values)
    # stable=True if there may be repeated elements in values
    # and you want always first occurrence
    idx_s = tf.argsort(values, stable=True)
    values_s = tf.gather(values, idx_s)
    idx_search = tf.searchsorted(values_s, target_values)
    return tf.gather(idx_s, idx_search)

print(find_in_array(values, target_values).numpy())
# [2 1]
user_17827920
Related