tensor 使用索引切片后的维度变化

import torch
a= torch.rand(10,5)
b=torch.tensor([1,4])
c=torch.tensor([2])
print(a[b].shape,a[c].shape, a[2].shape)

为什么a[2]的shape比a[c]的shape降低了维度?

在 PyTorch 中，数组的索引可以根据索引的类型影响返回的张量的形状。解释你的例子中不同索引结果的形状，可以从如何使用整数索引和张量索引进行理解。

1. a[b].shape：这里 b 是一个包含两个元素的张量（torch.tensor([1, 4])），所以 a[b] 会从 a 中选取索引为 1 和 4 的行。由于 b 是一维的，所以选出的结果仍然是二维的，具体来说是一个 2x5 的张量。

2. a[c].shape：这里 c 也是一个张量，但只包含一个元素（torch.tensor([2])）。根据 PyTorch 的机制，使用张量作为索引会保持原始张量的维度，所以 a[c] 返回的是一个 1x5 的张量。

3. a[2].shape：这里直接使用了一个整数进行索引，这在 PyTorch 中意味着你要访问第 3 行的元素（索引从 0 开始）。当你用一个单独的整数索引一个多维张量时，它会返回这个特定索引下的元素，而不是作为一个单独的维度。因此，它直接返回了一个向量，形状是 5（一维的）。

总结，当你使用单个整数索引时，PyTorch 会减少一个维度，因为你是在请求具体的行数据。当使用包含整数的张量进行索引时，它会保持原数组的维度，因为这被视为通过一个包含索引的列表来选择行，即使列表只有一个元素。这种差异是基于如何解释索引——是作为一个标量还是作为一个数组。

那d=torch.tensor([[1,4]])，a[d].shape是多少?

当你使用 d = torch.tensor([[1, 4]]) 这样的二维张量对 a 进行索引时，结果会有些不同。这里 d 是一个 1x2 的二维张量。

在 PyTorch 中，当你用一个二维张量进行索引时，会根据这个索引张量的每个元素从被索引张量中选择对应的行，然后将这些行根据索引张量的形状组合起来。因此，每个索引元素会抽取出一个行向量，然后这些向量按照索引张量的形状排列。

在这个例子中，d 的形状是 1x2，它包含了两个索引 1 和 4。因此，a[d] 将从 a 中抽取第 2 行和第 5 行（索引从 0 开始计数），并且因为索引是二维的，所以这两行会被组织为一个 1x2x5 的张量。

简单来说，a[d].shape 应该是 (1, 2, 5)。这里 1 对应于 d 的第一个维度，2 对应于 d 的第二个维度，5 是从 a 中每行选择出的元素数量。

还有另一种情况，比如在POMO的代码中，

self.BATCH_IDX = torch.arange(self.batch_size)[:, None].expand(self.batch_size, self.pomo_size)
self.POMO_IDX = torch.arange(self.pomo_size)[None, :].expand(self.batch_size, self.pomo_size)

self.current_node = selected
# shape: (batch, pomo)
self.step_state.ninf_mask[self.BATCH_IDX, self.POMO_IDX, self.current_node] = float('-inf')
# shape: (batch, pomo, node)