code: def partial_kernel_mmd_twostage(
      x,
      y,
      n_perm=None,
      kernel=GaussianKernel(),
      fraction_to_match=1.0,
      wasserstein_p=2.0,
  ):
      """Partial kernel MMD using a Wasserstein coupling to obtain the weight for the reference."""
      torchdrift.utils.check(
          n_perm is None,
          "Bootstrapping within partial MMD is not implemented, use bootstrap during fit",
          error_class=NotImplementedError,
      )
      n, d = x.shape
      m, d2 = y.shape
      if fraction_to_match < 1.0:
          _, coupling = torchdrift.detectors.wasserstein(
              x,
              y,
              fraction_to_match=fraction_to_match,
              return_coupling=True,
              n_perm=None,
              p=wasserstein_p,
          )
          w = (
              coupling[:, :-1].sum(1).to(device=x.device, dtype=x.dtype)
              / fraction_to_match
          )
      else:
          w = torch.full((n,), 1.0 / n, device=x.device, dtype=x.dtype)
  
      torchdrift.utils.check(d == d2, "feature dimension mismatch")
      xy = torch.cat([x.detach(), y.detach()], dim=0)
      dists = torch.cdist(xy, xy, p=2.0)
      # we are a bit sloppy here as we just keep the diagonal and everything twice
      k = kernel(dists)
      k_x = k[:n, :n]
      k_y = k[n:, n:]
      k_xy = k[:n, n:]
      # The diagonals are always 1 (up to numerical error, this is (3) in Gretton et al.)
      mmd = (
          (w @ k_x) @ w + k_y.sum() / (m * m) - 2 * (w @ k_xy).sum() / m
      )
      return mmd
  
file_path: TorchDrift/torchdrift/detectors/partial_mmd.py
imports: from typing import Optional
  import torch
  from .mmd import GaussianKernel
  from .detector import Detector
  import torchdrift
  import warnings
label_desc: Use image node_id 1 for calling a global function with example usage: partial_kernel_mmd_twostage(x, y, n_perm, kernel, fraction_to_match, wasserstein_p) and returns: mmd