tensroflow按需排列张量，查找指定信息

问题描述：

示例输入数据：
shapes_sample = tf.constant([[1, 2], [3, 4], [5, 6], [7, 8]], dtype=tf.int32)
types_sample = tf.constant([1, 2, 2, 3], dtype=tf.int32)
location = tf.constant([[10, 20], [30, 40], [50, 60], [70, 80]], dtype=tf.int32)
types_index_sample = tf.constant([0, 1, 0, 0], dtype=tf.int32)
因为这些数据都需要根据给定顺序进行排序，假设如上，如何拿到types_sample=2，types_index_sample =1对应的location？

---

解决方案：

import tensorflow as tf

# 示例输入数据
shapes_sample = tf.constant([[1, 2], [3, 4], [5, 6], [7, 8]], dtype=tf.int32)
types_sample = tf.constant([1, 2, 2, 3], dtype=tf.int32)
location = tf.constant([[10, 20], [30, 40], [50, 60], [70, 80]], dtype=tf.int32)
types_index_sample = tf.constant([0, 1, 0, 0], dtype=tf.int32)

# 条件筛选
cond = tf.logical_and(tf.equal(types_sample, 2), tf.equal(types_index_sample, 1))
filtered_shape = tf.boolean_mask(shapes_sample, cond)[0]
filtered_location = tf.boolean_mask(location, cond)[0]

# 打印结果
with tf.Session() as sess:
    filtered_location_result = sess.run(filtered_location)
    filtered_shape_result = sess.run(filtered_shape)
    print(filtered_location_result)
    print(filtered_shape_result)

这段代码首先使用 tf.logical_and 函数创建一个条件，要求 types_sample 等于 2 并且 types_index_sample 等于 1。然后使用 tf.boolean_mask 函数根据这个条件从 location 张量中提取符合条件的形状信息和位置信息。