import torch


def masked_temperature_softmax(x:torch.Tensor, idx:int, T:float=0.01):
    """Temperature scaled softmax on torch.Tensor masked by indices `idx`"""
    return torch.exp(x[idx, :] / T) / torch.exp(x[idx, :] / T).sum()