docstr: Batched (across the permutations) MMD^2 computation for the original test statistic and the permutations.
  
  Parameters
  ----------
  x_all
      Concatenated reference and test instances.
  perms
      List with permutation vectors.
  m
      Number of reference instances.
  n
      Number of test instances.
  
  Returns
  -------
  MMD^2 statistic for the original and permuted reference and test sets.
docstr_tok: ["Batched","(","across","the","permutations",")","MMD^2","computation","for","the","original","test","statistic","and","the","permutations",".","Parameters","--","--","--","--","--","x_all","Concatenated","reference","and","test","instances",".","perms","List","with","permutation","vectors",".","m","Number","of","reference","instances",".","n","Number","of","test","instances",".","Returns","--","--","--","-","MMD^2","statistic","for","the","original","and","permuted","reference","and","test","sets","."]
code: def _mmd2(
      self,
      x_all: torch.Tensor,
      perms: List[torch.Tensor],
      m: int,
      n: int,
  ) -> Tuple[torch.Tensor, torch.Tensor]:
      """
      Batched (across the permutations) MMD^2 computation for the original test statistic and the permutations.
  
      Parameters
      ----------
      x_all
          Concatenated reference and test instances.
      perms
          List with permutation vectors.
      m
          Number of reference instances.
      n
          Number of test instances.
  
      Returns
      -------
      MMD^2 statistic for the original and permuted reference and test sets.
      """
      k_xx, k_yy, k_xy = [], [], []
      for batch in range(self.n_batches):
          i, j = batch * self.batch_size, (batch + 1) * self.batch_size
          # construct stacked tensors with a batch of permutations for the reference set x and test set y
          x = torch.cat(
              [x_all[perm[:m]][None, :, :] for perm in perms[i:j]], 0
          )
          y = torch.cat(
              [x_all[perm[m:]][None, :, :] for perm in perms[i:j]], 0
          )
          if batch == 0:
              x = torch.cat([x_all[None, :m, :], x], 0)
              y = torch.cat([x_all[None, m:, :], y], 0)
          x, y = x.to(self.device), y.to(self.device)
  
          # batch-wise kernel matrix computation over the permutations
          k_xy.append(
              self.kernel(
                  LazyTensor(x[:, :, None, :]),
                  LazyTensor(y[:, None, :, :]),
                  self.infer_sigma,
              )
              .sum(1)
              .sum(1)
              .squeeze(-1)
          )
          k_xx.append(
              self.kernel(
                  LazyTensor(x[:, :, None, :]),
                  LazyTensor(x[:, None, :, :]),
              )
              .sum(1)
              .sum(1)
              .squeeze(-1)
          )
          k_yy.append(
              self.kernel(
                  LazyTensor(y[:, :, None, :]),
                  LazyTensor(y[:, None, :, :]),
              )
              .sum(1)
              .sum(1)
              .squeeze(-1)
          )
      c_xx, c_yy, c_xy = (
          1 / (m * (m - 1)),
          1 / (n * (n - 1)),
          2.0 / (m * n),
      )
      # Note that the MMD^2 estimates assume that the diagonal of the kernel matrix consists of 1's
      stats = (
          c_xx * (torch.cat(k_xx) - m)
          + c_yy * (torch.cat(k_yy) - n)
          - c_xy * torch.cat(k_xy)
      )
      return stats[0], stats[1:]
  
file_path: alibi-detect/alibi_detect/cd/keops/mmd.py
imports: import logging
  import numpy
  from pykeops.torch import LazyTensor
  import torch
  from typing import Callable, Dict, List, Optional, Tuple, Union
  from alibi_detect.cd.base import BaseMMDDrift
  from alibi_detect.utils.keops.kernels import GaussianRBF
  from alibi_detect.utils.pytorch import get_device
  from alibi_detect.utils.frameworks import Framework
  from alibi_detect.utils._types import TorchDeviceType
label_desc: Use image node_id 2 for calling the MMDDriftKeops obj's underlying member method code with example usage: obj._mmd2(x_all, perms, m, n) and returns: stats,  stats