支持 XLA 的动态切片答案

【问题标题】：XLA-able dynamic slicing支持 XLA 的动态切片
【发布时间】：2020-05-29 11:04:46
【问题描述】：

有没有办法根据 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 的参数在编译时已知。有推荐的解决方法吗？

【问题讨论】：

标签： python tensorflow tensorflow-xla

【解决方案1】：

这里的根本问题是 XLA 需要静态地知道程序中所有Tensors 的形状。在这种情况下，它抱怨tf.range，因为它的输出在随机输入的情况下是不可知的。相反，您可能能够摆脱生成一个屏蔽版本（将不需要的元素归零，使用类似 tensor_scatter_nd_update 之类的东西）并在下游使用该屏蔽版本（很难确切地说出如何，没有看到更多关于如何@987654323 @ 将被使用）。

【讨论】：