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: # Parameters
  
  predicted_labels : `torch.Tensor`, required.
      A tensor of predicted integer class labels of shape (batch_size, ...). Represented as C.
  gold_labels : `torch.Tensor`, required.
      A tensor of ground-truth integer class labels of shape (batch_size, ...). It must be the same
      shape as the `predicted_labels` tensor. Represented as Y.
  protected_variable_labels : `torch.Tensor`, required.
      A tensor of integer protected variable labels of shape (batch_size, ...). It must be the same
      shape as the `predicted_labels` tensor. Represented as A.
  mask : `torch.BoolTensor`, optional (default = `None`).
      A tensor of the same shape as `predicted_labels`.
  
  !!! Note
      All tensors are expected to be on the same device.
docstr_tok: ["#","Parameters","predicted_labels",":","`","torch.Tensor","`",",","required",".","A","tensor","of","predicted","integer","class","labels","of","shape","(","batch_size",",","...",")",".","Represented","as","C.","gold_labels",":","`","torch.Tensor","`",",","required",".","A","tensor","of","ground-truth","integer","class","labels","of","shape","(","batch_size",",","...",")",".","It","must","be","the","same","shape","as","the","`","predicted_labels","`","tensor",".","Represented","as","Y.","protected_variable_labels",":","`","torch.Tensor","`",",","required",".","A","tensor","of","integer","protected","variable","labels","of","shape","(","batch_size",",","...",")",".","It","must","be","the","same","shape","as","the","`","predicted_labels","`","tensor",".","Represented","as","A.","mask",":","`","torch.BoolTensor","`",",","optional","(","default","=","`","None","`",")",".","A","tensor","of","the","same","shape","as","`","predicted_labels","`",".","!","!","!","Note","All","tensors","are","expected","to","be","on","the","same","device","."]
code: def __call__(  # type: ignore
      self,
      predicted_labels: torch.Tensor,
      gold_labels: torch.Tensor,
      protected_variable_labels: torch.Tensor,
      mask: Optional[torch.BoolTensor] = None,
  ) -> None:
      """
      # Parameters
  
      predicted_labels : `torch.Tensor`, required.
          A tensor of predicted integer class labels of shape (batch_size, ...). Represented as C.
      gold_labels : `torch.Tensor`, required.
          A tensor of ground-truth integer class labels of shape (batch_size, ...). It must be the same
          shape as the `predicted_labels` tensor. Represented as Y.
      protected_variable_labels : `torch.Tensor`, required.
          A tensor of integer protected variable labels of shape (batch_size, ...). It must be the same
          shape as the `predicted_labels` tensor. Represented as A.
      mask : `torch.BoolTensor`, optional (default = `None`).
          A tensor of the same shape as `predicted_labels`.
  
      !!! Note
          All tensors are expected to be on the same device.
      """
      (
          predicted_labels,
          gold_labels,
          protected_variable_labels,
          mask,
      ) = self.detach_tensors(
          predicted_labels, gold_labels, protected_variable_labels, mask
      )
  
      # Some sanity checks.
      if predicted_labels.size() != protected_variable_labels.size():
          raise ConfigurationError(
              "protected_variable_labels must be of same size as predicted_labels but "
              "found tensor of shape: {}".format(
                  protected_variable_labels.size()
              )
          )
      if predicted_labels.size() != gold_labels.size():
          raise ConfigurationError(
              "gold_labels must be of same size as predicted_labels but "
              "found tensor of shape: {}".format(gold_labels.size())
          )
      if mask is not None and predicted_labels.size() != mask.size():
          raise ConfigurationError(
              "mask must be of same size as predicted_labels but "
              "found tensor of shape: {}".format(mask.size())
          )
      if (predicted_labels >= self._num_classes).any():
          raise ConfigurationError(
              "predicted_labels contains an id >= {}, "
              "the number of classes.".format(self._num_classes)
          )
      if (gold_labels >= self._num_classes).any():
          raise ConfigurationError(
              "gold_labels contains an id >= {}, "
              "the number of classes.".format(self._num_classes)
          )
      if (
          protected_variable_labels
          >= self._num_protected_variable_labels
      ).any():
          raise ConfigurationError(
              "protected_variable_labels contains an id >= {}, "
              "the number of protected variable labels.".format(
                  self._num_protected_variable_labels
              )
          )
  
      device = predicted_labels.device
      self._predicted_label_counts_by_gold_label = (
          self._predicted_label_counts_by_gold_label.to(device)
      )
      self._predicted_label_counts_by_gold_label_and_protected_variable_label = self._predicted_label_counts_by_gold_label_and_protected_variable_label.to(
          device
      )
      self._total_predictions = self._total_predictions.to(device)
  
      if mask is not None:
          predicted_labels = predicted_labels[mask]
          gold_labels = gold_labels[mask]
          protected_variable_labels = protected_variable_labels[mask]
      else:
          predicted_labels = predicted_labels.flatten()
          gold_labels = gold_labels.flatten()
          protected_variable_labels = (
              protected_variable_labels.flatten()
          )
  
      _total_predictions = torch.tensor(predicted_labels.nelement()).to(
          device
      )
      _predicted_label_counts_by_gold_label = torch.zeros(
          (self._num_classes, self._num_classes)
      ).to(device)
      _predicted_label_counts_by_gold_label_and_protected_variable_label = torch.zeros(
          (
              self._num_classes,
              self._num_protected_variable_labels,
              self._num_classes,
          )
      ).to(
          device
      )
      for y in range(self._num_classes):
          _predicted_label_counts_by_gold_label[y] = (
              predicted_labels[gold_labels == y]
              .float()
              .histc(
                  bins=self._num_classes,
                  min=0,
                  max=self._num_classes - 1,
              )
          )
          for a in range(self._num_protected_variable_labels):
              _predicted_label_counts_by_gold_label_and_protected_variable_label[
                  y
              ][
                  a
              ] = (
                  predicted_labels[
                      (gold_labels == y)
                      & (protected_variable_labels == a)
                  ]
                  .float()
                  .histc(
                      bins=self._num_classes,
                      min=0,
                      max=self._num_classes - 1,
                  )
              )
  
      if is_distributed():
          _total_predictions = _total_predictions.to(device)
          dist.all_reduce(_total_predictions, op=dist.ReduceOp.SUM)
  
          _predicted_label_counts_by_gold_label = (
              _predicted_label_counts_by_gold_label.to(device)
          )
          dist.all_reduce(
              _predicted_label_counts_by_gold_label[y],
              op=dist.ReduceOp.SUM,
          )
  
          _predicted_label_counts_by_gold_label_and_protected_variable_label = _predicted_label_counts_by_gold_label_and_protected_variable_label.to(
              device
          )
          dist.all_reduce(
              _predicted_label_counts_by_gold_label_and_protected_variable_label,
              op=dist.ReduceOp.SUM,
          )
  
      self._total_predictions += _total_predictions
      self._predicted_label_counts_by_gold_label += (
          _predicted_label_counts_by_gold_label
      )
      self._predicted_label_counts_by_gold_label_and_protected_variable_label += _predicted_label_counts_by_gold_label_and_protected_variable_label
  
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 2 for calling the Separation obj's underlying member method code with example usage: obj.__call__(predicted_labels, gold_labels, protected_variable_labels, mask) without return types