import torch


def argSoftmax2d(tensor, temperature, num_samples=1, return_1d_idx=False):
    """
  Find the index of the Softmax value in a 2d tensor. Prob is proportional to q/temperature

  Args:
    - tensor: PyTorch tensor of size (n x 1 x d x d)

  Returns: n x num_samples x 2 PyTorch tensor containing indexes of max values
  """
    n = tensor.size(0)
    d = tensor.size(2)
    probs = (tensor / temperature).view(n, -1).softmax(dim=-1)
    m = torch.multinomial(probs, num_samples)
    if not return_1d_idx:
        return torch.cat(((m // d).view(-1, 1), (m % d).view(-1, 1)), dim=1).long()
    else:
        return torch.cat(((m // d).view(-1, 1), (m % d).view(-1, 1)), dim=1).long(), m.long()