code: def partial_kernel_mmd_approx(
      x, y, fraction_to_match=1.0, kernel=GaussianKernel(), n_perm=None
  ):
      torchdrift.utils.check(
          n_perm is None,
          "Bootstrapping within partial MMD is not implemented, use bootstrap during fit",
          error_class=NotImplementedError,
      )
      rng = torch.Generator(device=x.device).manual_seed(1234)
      n, d = x.shape
      m, d2 = y.shape
      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)
      k = kernel(dists.double())
      k_x = k[:n, :n]
      k_y = k[n:, n:]
      k_xy = k[:n, n:]
  
      w = torch.full(
          (n,),
          1.0 / n,
          dtype=k_x.dtype,
          device=k_x.device,
          requires_grad=False,
      )
      mmd = (
          (w @ k_x) @ w + k_y.sum() / (m * m) - 2 * (w @ k_xy).sum() / m
      )
  
      for i in range(100):
          r = (
              torch.rand(
                  (), device=x.device, dtype=x.dtype, generator=rng
              )
              + 0.5
          )
          grad_mmd = (k_x @ w) - (k_xy).mean(1)
          grad_mmd_min = grad_mmd[(w < 1.0 / (n * fraction_to_match))]
          if grad_mmd_min.size(0) > 0:
              grad_mmd_min = grad_mmd_min.min()
          else:
              grad_mmd_min = torch.zeros_like(r)
          grad_mmd_max = grad_mmd[(w > 0)]
          if grad_mmd_max.size(0) > 0:
              grad_mmd_max = grad_mmd_max.max()
          else:  # pragma: no cover
              grad_mmd_max = torch.zeros_like(r)
          active_mask = ((w > 0) | (grad_mmd < grad_mmd_min * r)) & (
              (w < 1.0 / (n * fraction_to_match))
              | (grad_mmd > grad_mmd_max * r)
          )
  
          H_mmd_active = k_x[active_mask][:, active_mask]
          if H_mmd_active.size(0) == 0:
              continue
          u = torch.cholesky(H_mmd_active)
          Hinvgr = torch.cholesky_solve(
              grad_mmd[active_mask][:, None], u
          ).squeeze(1)
  
          w_active = w[active_mask]
  
          Hinvgr -= Hinvgr.mean()
          Hinvgr_full = torch.zeros_like(w)
          Hinvgr_full[active_mask] = Hinvgr
  
          step = 1.0
          for j in range(5):
              w_cand = w.clone()
              w_cand -= step * Hinvgr_full
  
              w_cand.clamp_(min=0, max=1.0 / (n * fraction_to_match))
              w_cand /= w_cand.sum()
              mmd_cand = (
                  (w_cand @ k_x) @ w_cand
                  + k_y.sum() / (m * m)
                  - 2 * (w_cand @ k_xy).sum() / m
              )
              is_lower = mmd_cand < mmd
              mmd = torch.where(is_lower, mmd_cand, mmd)
              w = torch.where(is_lower, w_cand, w)
              step /= 5
  
          grad_mmd = 2 * (k_x @ w) - 2 * (k_xy).mean(1)
          grad_mmd_min = grad_mmd[(w < 1.0 / (n * fraction_to_match))]
          if grad_mmd_min.size(0) > 0:
              grad_mmd_min = grad_mmd_min.min()
          else:
              grad_mmd_min = torch.zeros_like(r)
          grad_mmd_max = grad_mmd[(w > 0)]
          if grad_mmd_max.size(0) > 0:
              grad_mmd_max = grad_mmd_max.max()
          else:  # pragma: no cover
              grad_mmd_max = torch.zeros_like(r)
          active_mask = ((w > 0) | (grad_mmd < grad_mmd_min * r)) & (
              (w < 1.0 / (n * fraction_to_match))
              | (grad_mmd > grad_mmd_max * r)
          )
          step = 1e-1
          for j in range(5):
              w_candnd = w.clone()
              grad_mmd_x = grad_mmd.clone()
              grad_mmd_x = torch.where(
                  active_mask,
                  grad_mmd_x,
                  torch.zeros(
                      (),
                      device=grad_mmd_x.device,
                      dtype=grad_mmd_x.dtype,
                  ),
              )
              grad_mmd_x = torch.where(
                  active_mask,
                  grad_mmd_x,
                  grad_mmd_x - grad_mmd_x.mean(),
              )
              w_cand -= step * grad_mmd_x
              w_cand.clamp_(min=0, max=1.0 / (n * fraction_to_match))
              w_cand /= w_cand.sum()
              mmd_cand = (
                  (w_cand @ k_x) @ w_cand
                  + k_y.sum() / (m * m)
                  - 2 * (w_cand @ k_xy).sum() / m
              )
              is_lower = mmd_cand < mmd
              mmd = torch.where(is_lower, mmd_cand, mmd)
              w = torch.where(is_lower, w_cand, w)
              step = step / 5
  
      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 3 for calling a global function with example usage: partial_kernel_mmd_approx(x, y, fraction_to_match, kernel, n_perm) and returns: mmd