Python tf.dynamic_划分的PyTorch等价物

我需要在PyTorch中使用一个等价的

tf.dynamic_分区

。PyTorch或其他库中是否有类似的功能，或者是否有一种简单而聪明的方法来为PyTorch编码并快速工作？

是的，您可以为唯一索引迭代创建布尔掩码：

import torch


# partitions should be integer-like types
def dynamic_partition(
    data: torch.Tensor, partitions: torch.Tensor, num_partitions=None
):
    assert len(partitions.shape) == 1, "Only one dimensional partitions supported"
    assert (
        data.shape[0] == partitions.shape[0]
    ), "Partitions requires the same size as data"

    if num_partitions is None:
        num_partitions = max(torch.unique(partitions))

    return [data[partitions == i] for i in range(num_partitions)]

以及从文件中获取的示例：

partitions = torch.tensor([0, 0, 1, 1, 0])
num_partitions = 2
data = torch.tensor([10, 20, 30, 40, 50])

dynamic_partition(data, partitions, num_partitions)

其结果与示例相同，即：

[torch.tensor([10, 20, 50]), torch.tensor([30, 40])]