class_docstr: [Separation](https://fairmlbook.org) (pg. 12) allows correlation between the
  predictions and the protected variable to the extent that it is justified by
  the gold labels.
  
  # 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 predicted labels given a gold label and a sensitive attribute from the
      distribution over predicted labels given only the gold 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: ["[","Separation","]","(","https",":","\/\/fairmlbook.org",")","(","pg",".","12",")","allows","correlation","between","the","predictions","and","the","protected","variable","to","the","extent","that","it","is","justified","by","the","gold","labels",".","#","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","predicted","labels","given","a","gold","label","and","a","sensitive","attribute","from","the","distribution","over","predicted","labels","given","only","the","gold","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 y to a dictionary mapping each protected
      variable label a to the KL divergence or Wasserstein distance of P(C | A = a, Y = y) from
      P(C | Y = y). A distance of nearly 0 implies fairness on the basis of Separation.
  
  !!! Note
      If a class label is not present in Y 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","y","to","a","dictionary","mapping","each","protected","variable","label","a","to","the","KL","divergence","or","Wasserstein","distance","of","P","(","C","|","A","=","a",",","Y","=","y",")","from","P","(","C","|","Y","=","y",")",".","A","distance","of","nearly","0","implies","fairness","on","the","basis","of","Separation",".","!","!","!","Note","If","a","class","label","is","not","present","in","Y","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 y to a dictionary mapping each protected
          variable label a to the KL divergence or Wasserstein distance of P(C | A = a, Y = y) from
          P(C | Y = y). A distance of nearly 0 implies fairness on the basis of Separation.
  
      !!! Note
          If a class label is not present in Y 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 = {
              y: {
                  a: torch.tensor(float("nan"))
                  for a in range(self._num_protected_variable_labels)
              }
              for y in range(self._num_classes)
          }
          return distances
  
      for y in range(self._num_classes):
          probs = (
              self._predicted_label_counts_by_gold_label[y]
              / self._total_predictions
          )
          C_given_y_dist = Categorical(probs)
          if self._dist_metric == wasserstein_distance:
              C_given_y_dist = C_given_y_dist.probs
          distances[y] = {}
          for a in range(self._num_protected_variable_labels):
              probs = (
                  self._predicted_label_counts_by_gold_label_and_protected_variable_label[
                      y
                  ][
                      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[y][a] = torch.tensor(float("nan"))
                      continue
                  C_given_a_and_y_dist = Categorical(probs)
                  distances[y][a] = self._dist_metric(
                      C_given_a_and_y_dist, C_given_y_dist
                  )
              elif self._dist_metric == wasserstein_distance:
                  C_given_a_and_y_dist = Categorical(probs).probs
                  # scipy expects empirical values as well as probabilities
                  label_values = torch.tensor(range(self._num_classes))
                  distances[y][a] = self._dist_metric(
                      label_values,
                      label_values,
                      C_given_a_and_y_dist,
                      C_given_y_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 Separation obj's underlying member method code with example usage: obj.get_metric(reset) and returns: distances, distances