Python 在TensorFlow中，如何断言列表的值在特定集合中？

我有一个一维的

tf.uint8

张量

x

，我想断言该张量内的所有值都在我定义的集合

s

中

s

在图形定义时是固定的，因此它不是动态计算的张量

在普通Python中，我想做如下事情：

x = [1, 2, 3, 1, 11, 3, 5]
s = {1, 2, 3, 11, 12, 13}
assert all(el in s for el in x), "This should fail, as 5 is not in s"

我知道我可以使用断言部分，但我正在努力定义条件部分（

el in s

）。最简单/最规范的方法是什么

---

旧的答案对我来说是不够的：首先，写下来和理解起来很复杂，其次，它使用的是广播的

tf.equal

，这在计算方面比正确的基于集合的检查更昂贵。

一个简单的方法可能是这样的：

import tensorflow as tf

x = [1, 2, 3, 1, 11, 3, 5]
s = {1, 2, 3, 11, 12, 13}

x_t = tf.constant(x, dtype=tf.uint8)
s_t = tf.constant(list(s), dtype=tf.uint8)
# Check every value in x against every value in s
xs_eq = tf.equal(x_t[:, tf.newaxis], s_t)
# Check every element in x is equal to at least one element in s
assert_op = tf.Assert(tf.reduce_all(tf.reduce_any(xs_eq, axis=1)), [x_t])
with tf.control_dependencies([assert_op]):
    # Use x_t...

import tensorflow as tf

x = [1, 2, 3, 1, 11, 3, 5]
s = {1, 2, 3, 11, 12, 13}

x_t = tf.constant(x, dtype=tf.uint8)
# Count where each x matches each s
x_in_s = [tf.cast(tf.equal(x_t, si), tf.int32) for si in s]
# Add matches and check there is at least one match per x
assert_op = tf.Assert(tf.reduce_all(tf.add_n(x_in_s) > 0), [x_t])

这将创建一个大小为

（len（x），len（s））

的中间张量。如果有问题，您还可以将问题拆分为独立的张量，例如：

import tensorflow as tf

x = [1, 2, 3, 1, 11, 3, 5]
s = {1, 2, 3, 11, 12, 13}

x_t = tf.constant(x, dtype=tf.uint8)
s_t = tf.constant(list(s), dtype=tf.uint8)
# Check every value in x against every value in s
xs_eq = tf.equal(x_t[:, tf.newaxis], s_t)
# Check every element in x is equal to at least one element in s
assert_op = tf.Assert(tf.reduce_all(tf.reduce_any(xs_eq, axis=1)), [x_t])
with tf.control_dependencies([assert_op]):
    # Use x_t...

import tensorflow as tf

x = [1, 2, 3, 1, 11, 3, 5]
s = {1, 2, 3, 11, 12, 13}

x_t = tf.constant(x, dtype=tf.uint8)
# Count where each x matches each s
x_in_s = [tf.cast(tf.equal(x_t, si), tf.int32) for si in s]
# Add matches and check there is at least one match per x
assert_op = tf.Assert(tf.reduce_all(tf.add_n(x_in_s) > 0), [x_t])

编辑：

实际上，因为您说过您的值是

tf.uint8

，所以使用布尔数组可以让事情变得更好：

import tensorflow as tf

x = [1, 2, 3, 1, 11, 3, 5]
s = {1, 2, 3, 11, 12, 13}

x_t = tf.constant(x, dtype=tf.uint8)
s_t = tf.constant(list(s), dtype=tf.uint8)
# One-hot vectors of values included in x and s
x_bool = tf.scatter_nd(tf.cast(x_t[:, tf.newaxis], tf.int32),
                       tf.ones_like(x_t, dtype=tf.bool), [256])
s_bool = tf.scatter_nd(tf.cast(s_t[:, tf.newaxis], tf.int32),
                       tf.ones_like(s_t, dtype=tf.bool), [256])
# Check that all values in x are in s
assert_op = tf.Assert(tf.reduce_all(tf.equal(x_bool, x_bool & s_bool)), [x_t])

这需要线性时间和恒定内存

编辑2：虽然最后一种方法在理论上是这种情况下最好的，但在进行两次快速基准测试时，我只能看到当我使用数十万个元素时，性能上的显著差异，而且在任何情况下，这三种方法在

tf.uint8

中仍然相当快