class_docstr: [Sufficiency](https://fairmlbook.org) (pg. 14) is satisfied by the predictions
  when the protected variable and gold labels are clear from context.
  
  # Parameters
  
  num_classes : `int`
      Number of classes.
  num_protected_variable_labels : `int`
      Number of protected variable labels.
  dist_metric : `str`
      Distance metric (kl_divergence, wasserstein) for calculating the distance between
      the distribution over gold labels given a predicted label and a sensitive attribute from the
      distribution over gold labels given only the predicted label. If both distributions do not
      have equal support, you should use wasserstein distance.
  
  !!! Note
      Assumes integer labels, with each item to be classified having
      a single correct class.
class_docstr_tok: ["[","Sufficiency","]","(","https",":","\/\/fairmlbook.org",")","(","pg",".","14",")","is","satisfied","by","the","predictions","when","the","protected","variable","and","gold","labels","are","clear","from","context",".","#","Parameters","num_classes",":","`","int","`","Number","of","classes",".","num_protected_variable_labels",":","`","int","`","Number","of","protected","variable","labels",".","dist_metric",":","`","str","`","Distance","metric","(","kl_divergence",",","wasserstein",")","for","calculating","the","distance","between","the","distribution","over","gold","labels","given","a","predicted","label","and","a","sensitive","attribute","from","the","distribution","over","gold","labels","given","only","the","predicted","label",".","If","both","distributions","do","not","have","equal","support",",","you","should","use","wasserstein","distance",".","!","!","!","Note","Assumes","integer","labels",",","with","each","item","to","be","classified","having","a","single","correct","class","."]
docstr: # 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.
docstr_tok: ["#","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","."]
code: 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
  
imports: from typing import Optional, Dict
  from scipy.stats import wasserstein_distance
  import torch
  import torch.distributed
  from torch.distributions.categorical import Categorical
  from torch.distributions.kl import kl_divergence
  from allennlp.common.util import is_distributed
  from allennlp.common.checks import ConfigurationError
  from allennlp.training.metrics.metric import Metric
label_desc: Use image node_id 3 for calling the Sufficiency obj's underlying member method code with example usage: obj.get_metric(reset) and returns: distances, distances