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: Union[list, 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.
      """
      preprocess_batch_fn = self.train_kwargs["preprocess_fn"]
      if isinstance(preprocess_batch_fn, Callable):  # type: ignore[arg-type]
          x_all = preprocess_batch_fn(x_all)  # type: ignore[operator]
      if self.has_proj:
          x_all_proj = predict_batch(
              x_all,
              self.kernel.proj,
              device=self.device,
              batch_size=self.batch_size_predict,
              dtype=x_all.dtype
              if isinstance(x_all, torch.Tensor)
              else torch.float32,
          )
      else:
          x_all_proj = x_all
  
      x, x2, y, y2 = None, None, None, None
      k_xx, k_yy, k_xy = [], [], []
      for batch in range(self.n_batches):
          i, j = (
              batch * self.batch_size_perms,
              (batch + 1) * self.batch_size_perms,
          )
          # Stack a batch of permuted reference and test tensors and their projections
          x_proj = torch.cat(
              [x_all_proj[perm[:m]][None, :, :] for perm in perms[i:j]],
              0,
          )
          y_proj = torch.cat(
              [x_all_proj[perm[m:]][None, :, :] for perm in perms[i:j]],
              0,
          )
          if self.has_kernel_b:
              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_proj = torch.cat([x_all_proj[None, :m, :], x_proj], 0)
              y_proj = torch.cat([x_all_proj[None, m:, :], y_proj], 0)
              if self.has_kernel_b:
                  x = torch.cat([x_all[None, :m, :], x], 0)  # type: ignore[call-overload]
                  y = torch.cat([x_all[None, m:, :], y], 0)  # type: ignore[call-overload]
          x_proj, y_proj = x_proj.to(self.device), y_proj.to(
              self.device
          )
          if self.has_kernel_b:
              x, y = x.to(self.device), y.to(self.device)
  
          # Batch-wise kernel matrix computation over the permutations
          with torch.no_grad():
              x2_proj, x_proj = LazyTensor(
                  x_proj[:, None, :, :]
              ), LazyTensor(x_proj[:, :, None, :])
              y2_proj, y_proj = LazyTensor(
                  y_proj[:, None, :, :]
              ), LazyTensor(y_proj[:, :, None, :])
              if self.has_kernel_b:
                  x2, x = LazyTensor(x[:, None, :, :]), LazyTensor(
                      x[:, :, None, :]
                  )
                  y2, y = LazyTensor(y[:, None, :, :]), LazyTensor(
                      y[:, :, None, :]
                  )
              k_xy.append(
                  self.kernel(x_proj, y2_proj, x, y2)
                  .sum(1)
                  .sum(1)
                  .squeeze(-1)
              )
              k_xx.append(
                  self.kernel(x_proj, x2_proj, x, x2)
                  .sum(1)
                  .sum(1)
                  .squeeze(-1)
              )
              k_yy.append(
                  self.kernel(y_proj, y2_proj, y, y2)
                  .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:]
  
imports: from copy import deepcopy
  from functools import partial
  from tqdm import tqdm
  import numpy
  from pykeops.torch import LazyTensor
  import torch
  import torch.nn
  from torch.utils.data import DataLoader
  from typing import Callable, Dict, List, Optional, Union, Tuple
  from alibi_detect.cd.base import BaseLearnedKernelDrift
  from alibi_detect.utils.pytorch import get_device, predict_batch
  from alibi_detect.utils.pytorch.data import TorchDataset
  from alibi_detect.utils.frameworks import Framework
  from alibi_detect.utils._types import TorchDeviceType
label_desc: Use image node_id 3 for calling the LearnedKernelDriftKeops obj's underlying member method code with example usage: obj._mmd2(x_all, perms, m, n) and returns: stats,  stats