code: def masked_pick_probability(x, y, temp, cos_distance):
      
      return SNNLCrossEntropy.pick_probability(x, temp, cos_distance) * \
          SNNLCrossEntropy.same_label_mask(y, y)
docs: The pairwise sampling probabilities for the elements of x for neighbor
      points which share labels.
      :param x: a matrix
      :param y: a list of labels for each element of x
      :param temp: Temperature
      :cos_distance: Boolean for using cosine or Euclidean distance
  
      :returns: A tensor for the pairwise sampling probabilities.