code: def wasserstein(
      x,
      y,
      p=2.0,
      fraction_to_match=1.0,
      n_perm=1000,
      return_coupling=False,
  ):
      """Wasserstein distance based two-sample test with partial match.
      (see T. Viehmann: Partial Wasserstein and Maximum Mean Discrepancy distances
      for bridging the gap between outlier detection and drift detection,
      https://arxiv.org/abs/2106.01289 )
  
      Args:
          x: reference distribution sample
          y: test distribution sample
          fraction_to_match: fraction of x probability mass to be matched
          n_perm: number of bootstrap permutations to run to compute p-value (None for not returning a p-value)
          return_coupling: return the coupling
  
      Note that this currently runs on the CPU and in 64 bit but it will
      move and cast as needed.
      """
      device = x.device
      dtype = x.dtype
      x = x.to(dtype=torch.float64, device="cpu")
      y = y.to(dtype=torch.float64, device="cpu")
      num_x, d = x.shape
      num_y, d2 = y.shape
      torchdrift.utils.check(
          d == d2, "Shape mismatch in feature dimension"
      )
  
      dists_p = torch.cdist(x, y, p=p) ** p
      if fraction_to_match < 1.0:
          max_dists_p = dists_p.max()
          dists_p = torch.cat(
              [dists_p, (1.1 * max_dists_p).expand(num_x, 1)], dim=1
          )
  
      weights_x = torch.full_like(dists_p[:, 0], 1.0 / num_x)[None]
      weights_y = torch.full_like(
          dists_p[0, :], fraction_to_match / num_y
      )[None]
      if fraction_to_match < 1.0:
          weights_y[:, -1] = 1.0 - fraction_to_match
  
      coupling = torch.from_numpy(
          ot_emd(
              weights_x[0].numpy(),
              weights_y[0].numpy(),
              dists_p.cpu().numpy(),
          )
      )
  
      if (
          coupling[:, :num_y].sum() / fraction_to_match - 1
      ).abs().item() > 1e-5:  # pragma: no cover
          raise RuntimeError("Numerical stability failed")
      wdist = (
          (coupling[:, :num_y] * dists_p[:, :num_y]).sum()
          / fraction_to_match
      ) ** (1 / p)
  
      if n_perm is None and return_coupling:
          return wdist.to(dtype=dtype, device=device), coupling.to(
              dtype=dtype, device=device
          )
      elif n_perm is None:
          return wdist.to(dtype=dtype, device=device)
  
      xy = torch.cat([x, y], dim=0)
      scores = []
      for i in range(n_perm):
          slicing = torch.randperm(num_x + num_y)
          dists_p_0 = (
              torch.cdist(xy[slicing[:num_x]], xy[slicing[num_x:]], p=p)
              ** p
          )
          if fraction_to_match < 1.0:
              max_dists_p_0 = dists_p_0.max()
              dists_p_0 = torch.cat(
                  [dists_p_0, (1.1 * max_dists_p_0).expand(num_x, 1)],
                  dim=1,
              )
          coupling_0 = torch.from_numpy(
              ot_emd(weights_x[0], weights_y[0], dists_p_0)
          )
  
          if (
              coupling_0[:, :num_y].sum() / fraction_to_match - 1
          ).abs().item() > 1e-5:  # pragma: no cover
              raise RuntimeError("Numerical stability failed")
          scores.append(
              (
                  (coupling_0[:, :num_y] * dists_p_0[:, :num_y]).sum()
                  / fraction_to_match
              )
              ** (1 / p)
          )
      scores = torch.stack(scores)
  
      p_val = (wdist < scores).float().mean()
      if return_coupling:
          return (
              wdist.to(dtype=dtype, device=device),
              p_val.to(dtype=dtype, device=device),
              coupling.to(dtype=dtype, device=device),
          )
      return wdist.to(dtype=dtype, device=device), p_val.to(
          dtype=dtype, device=device
      )
  
file_path: TorchDrift/torchdrift/detectors/wasserstein.py
imports: from typing import Optional
  import torch
  from .None import Detector
  import torchdrift.utils
label_desc: Use image node_id 1 for calling a global function with example usage: wasserstein(x, y, p, fraction_to_match, n_perm, return_coupling) and returns: wdist,  p_val, wdist,  coupling, wdist,  p_val,  coupling, wdist