TPU上的碎裂传感器

Question

我正在尝试使用TensorFlow来训练一个神经网络，该网络以aRaggedTensor作为输入。它在CPU和GPU上运行良好，但我真的在努力使它在TPU上运行。我想知道你们中的一些人是否能成功（不一定要寻找直接的解决方案，尽管这很好，但一些工具提示已经很好了！）。到目前为止，错误信息已经足够明确，我可以继续讲下去了，但我现在仍在为如何更进一步而挣扎

到目前为止我所做的：

1. 我正在使用tf.data.Dataset从tfu记录中读取数据，但我需要将其显式转换为DistributedDataset以禁用优化

strategy.experimental_distribute_dataset(
    dataset,
    tf.distribute.InputOptions(
        experimental_prefetch_to_device=False
    )
)

1. 我得到了Compilation failure: Detected unsupported operations when trying to compile graph ... on XLA_TPU_JIT: RaggedTensorToTensor，通过允许软设备放置，可以（某种程度上）修复它：

tf.config.set_soft_device_placement(True)

1. 我现在被Compilation failure: Input 1 to node '.../RaggedReduceSum/RaggedReduce/RaggedSplitsToSegmentIds/Repeat/SequenceMask/Range' with op Range must be a compile-time constant.困住了。我完全理解为什么会出现这个错误，我完全了解available ops on TPU，特别是大多数动态操作应该在编译时确定在TPU上运行。但我想不出我怎么能用这些不规则的张量和TPU

任何想法都将不胜感激：）

p.S.：从2020年7月开始，我就没有看到TensorFlow团队关于TPU上RaggedTensors的太多消息，但我可能错过了很多关于它的消息。。。如果我能进行更多的研究，那么指向git线程对我来说已经很好了