tensorflow output_data[output_id == i] = input_data[input_id == i]代码分析

import tensorflow as tf
import os
os.environ['CUDA_VISIBLE_DEVICES'] = '1'

with tf.Graph().as_default(), tf.Session() as sess:
    input_data = tf.constant([0.1, 0.2, 0.3])
    input_ids = tf.constant([3, 1, 6])
    output_data = tf.constant([0., 0., 0., 0., 0.])
    output_ids = tf.constant([6, 3, 1, 3, 0])
    # From TF v1.13
    #s = tf.argsort(input_ids)
    # Before TF v1.13
    s = tf.contrib.framework.argsort(input_ids)
    input_ids_s = tf.gather(input_ids, s)
    n = tf.size(input_ids)
    output_idx_s = tf.minimum(tf.searchsorted(input_ids_s, output_ids), n - 1)
    searchsorted = tf.searchsorted(input_ids_s, output_ids), n - 1
    gather1 = tf.gather(input_ids_s, output_idx_s)
    equal = tf.equal(output_ids, tf.gather(input_ids_s, output_idx_s))
    gather2 = tf.gather(s, output_idx_s)
    gather3 = tf.gather(input_data, tf.gather(s, output_idx_s))