    def get_metric(self, reset: bool = False) -> Dict[int, Dict[int, torch.FloatTensor]]:
        """
        # Returns

        distances : `Dict[int, Dict[int, torch.FloatTensor]]`
            A dictionary mapping each class label c to a dictionary mapping each protected
            variable label a to the KL divergence or Wasserstein distance of P(Y | A = a, C = c)
            from P(Y | C = c). A distance of nearly 0 implies fairness on the basis of Sufficiency.

        !!! Note
            If a possible class label is not present in C, the expected behavior is that
            the KL divergences corresponding to this class label are NaN. If a possible class label is
            not present in C conditioned on a protected variable label, the expected behavior is that
            the KL divergence corresponding to this (class label, protected variable label) pair is NaN.
            You can avoid this by using Wasserstein distance instead.
        """
        distances: Dict[int, Dict[int, torch.FloatTensor]] = {}
        if self._total_predictions == 0:
            distances = {
                c: {
                    a: torch.tensor(float("nan"))
                    for a in range(self._num_protected_variable_labels)
                }
                for c in range(self._num_classes)
            }
            return distances

        for c in range(self._num_classes):
            # It is possible that `c` is not predicted at all,
            # in which case `probs` is all zeros.
            probs = self._gold_label_counts_by_predicted_label[c] / self._total_predictions
            if self._dist_metric == kl_divergence:
                if probs.sum() == 0:
                    distances[c] = {
                        a: torch.tensor(float("nan"))
                        for a in range(self._num_protected_variable_labels)
                    }
                    continue
            Y_given_c_dist = Categorical(probs)
            distances[c] = {}
            if self._dist_metric == wasserstein_distance:
                Y_given_c_dist = Y_given_c_dist.probs
            for a in range(self._num_protected_variable_labels):
                probs = (
                    self._gold_label_counts_by_predicted_label_and_protected_variable_label[c][a]
                    / self._total_predictions
                )
                if self._dist_metric == kl_divergence:
                    # Implies class label y is not present in Y conditioned on protected variable label a
                    if probs.sum() == 0:
                        distances[c][a] = torch.tensor(float("nan"))
                        continue
                    Y_given_a_and_c_dist = Categorical(probs)
                    distances[c][a] = self._dist_metric(Y_given_a_and_c_dist, Y_given_c_dist)
                elif self._dist_metric == wasserstein_distance:
                    Y_given_a_and_c_dist = Categorical(probs).probs
                    # scipy expects empirical values as well as probabilities
                    label_values = torch.tensor(range(self._num_classes))
                    distances[c][a] = self._dist_metric(
                        label_values, label_values, Y_given_a_and_c_dist, Y_given_c_dist
                    )
        if reset:
            self.reset()
        return distances
# Returns

distances : `Dict[int, Dict[int, torch.FloatTensor]]`
A dictionary mapping each class label c to a dictionary mapping each protected
variable label a to the KL divergence or Wasserstein distance of P(Y | A = a, C = c)
from P(Y | C = c). A distance of nearly 0 implies fairness on the basis of Sufficiency.

!!! Note
If a possible class label is not present in C, the expected behavior is that
the KL divergences corresponding to this class label are NaN. If a possible class label is
not present in C conditioned on a protected variable label, the expected behavior is that
the KL divergence corresponding to this (class label, protected variable label) pair is NaN.
You can avoid this by using Wasserstein distance instead.
allennlp/fairness/fairness_metrics.py