动态形状

动态形状是指张量形状的可变性，其中其形状取决于另一个上游张量的值。例如：

>>> import torch, torch_xla
>>> in_tensor  = torch.randint(low=0, high=2, size=(5,5), device='xla:0')
>>> out_tensor = torch.nonzero(in_tensor)

out_tensor 的形状取决于 in_tensor 的值，并且受 in_tensor 的形状限制。换句话说，如果你执行

>>> print(out_tensor.shape)
torch.Size([<=25, 2])

可以看到第一维取决于 in_tensor 的值，并且其最大值为 25。我们将第一维称为动态维度。第二维不依赖于任何上游张量，因此我们称其为静态维度。

动态形状可进一步分为有界动态形状和无界动态形状。

• 有界动态形状：是指其动态维度受静态值限制的形状。它适用于需要静态内存分配的加速器（例如 TPU）。

• 无界动态形状：是指其动态维度可以无限大的形状。它适用于不需要静态内存分配的加速器（例如 GPU）。

目前，仅支持有界动态形状，并且尚处于实验阶段。

有界动态形状

目前，我们支持在 TPU 上具有动态输入尺寸的多层感知机模型 (MLP)。

此功能由标志 XLA_EXPERIMENTAL="nonzero:masked_select" 控制。要运行启用此功能Thus, we call the first dimension the dynamic dimension. The second dimension does not depend on any upstream tensors so we call it the static dimension.的模型，你可以执行

XLA_EXPERIMENTAL="nonzero:masked_select:masked_scatter" python your_scripts.py

以下是我们运行 MLP 模型 100 次迭代时获得的一些数字：

	无动态形状	有动态形状
端到端训练时间	29.49	20.03
编译次数	102	49
编译缓存命中率	198	1953

动态形状的动机之一是减少形状在迭代之间持续变化时过多的重新编译。从上图可以看出，编译次数减少了一半，从而降低了训练时间。

要尝试一下，请运行：

XLA_EXPERIMENTAL="nonzero:masked_select" PJRT_DEVICE=TPU python3 pytorch/xla/test/ds/test_dynamic_shape_models.py TestDynamicShapeModels.test_backward_pass_with_dynamic_input

有关我们计划将来在 PyTorch/XLA 上扩展动态形状支持的更多详细信息，请随意查阅我们的 RFC。