支持XLA的动态切片

Question

在XLA编译的函数中，有没有办法根据随机数生成器对张量进行动态切片？例如：

@tf.function(experimental_compile=True)
def random_slice(input, max_slice_size):
    offset = tf.squeeze(tf.random.uniform([1], minval=0, maxval=input.shape[0]-max_slice_size, dtype=tf.int32))
    sz = tf.squeeze(tf.random.uniform([1], minval=1, maxval=max_slice_size, dtype=tf.int32))

    indices = tf.range(offset, offset+sz)  # Non-XLA-able due to non-static bounds

    return tf.gather(input, indices)

x = tf.ones([50, 50])
y = random_slice(x, 4)

这段代码无法编译，因为XLA要求tf.range的参数在编译时是已知的。是否有推荐的解决方法？

Answer 1

这里的基本问题是，XLA需要静态地知道程序中所有Tensor的形状。在这种情况下，它抱怨tf.range，因为在给定随机输入的情况下，它的输出是未知的。相反，您可能能够生成一个屏蔽版本(将不需要的元素清零，使用类似tensor_scatter_nd_update的东西)，并在下游使用该屏蔽版本(很难确切地说出如何使用，因为没有看到关于如何使用y的更多上下文)。