python - 具有大张量的更快的torch.quantile？

Question

我有一个非常大的张量，我想计算张量中每个元素的第 k 个百分位值。目前这是正在做的事情：

C = torch.quantile(A, B, dim=1, keepdim=True)
# Shape of A: (1000, 128, 128)
# Shape of B: (1000,)

在他们的文档中给出，quantile似乎对张量进行排序以计算百分位值。有更快的库支持快速排序，例如torchsort，我想知道我是否可以做这样的事情tensor => torchsort => quantile而quantile不必自己排序？

标签： pythonpytorch