class_docstr: Compute the generalized Wasserstein Dice Loss defined in:
  
      Fidon L. et al. (2017) Generalised Wasserstein Dice Score for Imbalanced Multi-class
      Segmentation using Holistic Convolutional Networks. BrainLes 2017.
  
  Or its variant (use the option weighting_mode="GDL") defined in the Appendix of:
  
      Tilborghs, S. et al. (2020) Comparative study of deep learning methods for the automatic
      segmentation of lung, lesion and lesion type in CT scans of COVID-19 patients.
      arXiv preprint arXiv:2007.15546
  
  Adapted from:
      https://github.com/LucasFidon/GeneralizedWassersteinDiceLoss
class_docstr_tok: ["Compute","the","generalized","Wasserstein","Dice","Loss","defined","in",":","Fidon","L.","et","al",".","(","2017",")","Generalised","Wasserstein","Dice","Score","for","Imbalanced","Multi-class","Segmentation","using","Holistic","Convolutional","Networks",".","BrainLes","2017",".","Or","its","variant","(","use","the","option","weighting_mode=","''","GDL","''",")","defined","in","the","Appendix","of",":","Tilborghs",",","S.","et","al",".","(","2020",")","Comparative","study","of","deep","learning","methods","for","the","automatic","segmentation","of","lung",",","lesion","and","lesion","type","in","CT","scans","of","COVID-19","patients",".","arXiv","preprint","arXiv:2007.15546","Adapted","from",":","https",":","\/\/github.com\/LucasFidon\/GeneralizedWassersteinDiceLoss"]
docstr: Compute the voxel-wise Wasserstein distance between the
  flattened prediction and the flattened labels (ground_truth) with respect
  to the distance matrix on the label space M.
  This corresponds to eq. 6 in:
  
      Fidon L. et al. (2017) Generalised Wasserstein Dice Score for Imbalanced Multi-class
      Segmentation using Holistic Convolutional Networks. BrainLes 2017.
  
  Args:
      flat_proba: the probabilities of input(predicted) tensor.
      flat_target: the target tensor.
docstr_tok: ["Compute","the","voxel-wise","Wasserstein","distance","between","the","flattened","prediction","and","the","flattened","labels","(","ground_truth",")","with","respect","to","the","distance","matrix","on","the","label","space","M.","This","corresponds","to","eq",".","6","in",":","Fidon","L.","et","al",".","(","2017",")","Generalised","Wasserstein","Dice","Score","for","Imbalanced","Multi-class","Segmentation","using","Holistic","Convolutional","Networks",".","BrainLes","2017",".","Args",":","flat_proba",":","the","probabilities","of","input","(","predicted",")","tensor",".","flat_target",":","the","target","tensor","."]
code: def wasserstein_distance_map(
      self, flat_proba: torch.Tensor, flat_target: torch.Tensor
  ) -> torch.Tensor:
      """
      Compute the voxel-wise Wasserstein distance between the
      flattened prediction and the flattened labels (ground_truth) with respect
      to the distance matrix on the label space M.
      This corresponds to eq. 6 in:
  
          Fidon L. et al. (2017) Generalised Wasserstein Dice Score for Imbalanced Multi-class
          Segmentation using Holistic Convolutional Networks. BrainLes 2017.
  
      Args:
          flat_proba: the probabilities of input(predicted) tensor.
          flat_target: the target tensor.
      """
      # Turn the distance matrix to a map of identical matrix
      m = torch.clone(torch.as_tensor(self.m)).to(flat_proba.device)
      m_extended = torch.unsqueeze(m, dim=0)
      m_extended = torch.unsqueeze(m_extended, dim=3)
      m_extended = m_extended.expand(
          (
              flat_proba.size(0),
              m_extended.size(1),
              m_extended.size(2),
              flat_proba.size(2),
          )
      )
  
      # Expand the feature dimensions of the target
      flat_target_extended = torch.unsqueeze(flat_target, dim=1)
      flat_target_extended = flat_target_extended.expand(
          (flat_target.size(0), m_extended.size(1), flat_target.size(1))
      )
      flat_target_extended = torch.unsqueeze(
          flat_target_extended, dim=1
      )
  
      # Extract the vector of class distances for the ground-truth label at each voxel
      m_extended = torch.gather(
          m_extended, dim=1, index=flat_target_extended
      )
      m_extended = torch.squeeze(m_extended, dim=1)
  
      # Compute the wasserstein distance map
      wasserstein_map = m_extended * flat_proba
  
      # Sum over the classes
      wasserstein_map = torch.sum(wasserstein_map, dim=1)
      return wasserstein_map
  
imports: from __future__ import annotations
  import warnings
  from collections.abc import Callable, Sequence
  from typing import Any
  import numpy
  import torch
  import torch.nn
  import torch.nn.functional
  from torch.nn.modules.loss import _Loss
  from monai.losses.focal_loss import FocalLoss
  from monai.losses.spatial_mask import MaskedLoss
  from monai.networks import one_hot
  from monai.utils import DiceCEReduction, LossReduction, Weight, deprecated_arg, look_up_option, pytorch_after
label_desc: Use image node_id 3 for calling the GeneralizedWassersteinDiceLoss obj's underlying member method code with example usage: obj.wasserstein_distance_map(flat_proba, flat_target) and returns: wasserstein_map