docstr: Private method to compute the MMD-ADiTT test statistic.
docstr_tok: ["Private","method","to","compute","the","MMD-ADiTT","test","statistic","."]
code: def _cmmd(
      self,
      K: torch.Tensor,
      L: torch.Tensor,
      bools: torch.Tensor,
      L_held: torch.Tensor = None,
  ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
      """
      Private method to compute the MMD-ADiTT test statistic.
      """
      # Get ref/test indices
      idx_0, idx_1 = (
          torch.where(bools == 0)[0],
          torch.where(bools == 1)[0],
      )
      n_ref, n_test = len(idx_0), len(idx_1)
  
      # Form kernel matrices
      L_0, L_1 = L[idx_0][:, idx_0], L[idx_1][:, idx_1]
      K_0, K_1 = K[idx_0][:, idx_0], K[idx_1][:, idx_1]
  
      # Initialise regularisation parameters
      # Implemented only for first _cmmd call which corresponds to original window assignment
      if self.lams is None:
          possible_lams = torch.tensor(
              [2 ** (-i) for i in range(20)]
          ).to(K.device)
          lam_0 = self._pick_lam(
              possible_lams, K_0, L_0, n_folds=self.n_folds
          )
          lam_1 = self._pick_lam(
              possible_lams, K_1, L_1, n_folds=self.n_folds
          )
          self.lams = (lam_0, lam_1)
  
      # Compute stat
      L_0_inv = torch.linalg.inv(
          L_0
          + n_ref * self.lams[0] * torch.eye(int(n_ref)).to(L_0.device)
      )
      L_1_inv = torch.linalg.inv(
          L_1
          + n_test
          * self.lams[1]
          * torch.eye(int(n_test)).to(L_1.device)
      )
      A_0 = L_held[:, idx_0] @ L_0_inv
      A_1 = L_held[:, idx_1] @ L_1_inv
      # Allow batches of MMDs to be computed at a time (rather than all)
      if self.batch_size is not None:
          bs = self.batch_size
          coupling_xx = torch.stack(
              [
                  torch.einsum("ij,ik->ijk", A_0_i, A_0_i).mean(0)
                  for A_0_i in A_0.split(bs)
              ]
          ).mean(0)
          coupling_yy = torch.stack(
              [
                  torch.einsum("ij,ik->ijk", A_1_i, A_1_i).mean(0)
                  for A_1_i in A_1.split(bs)
              ]
          ).mean(0)
          coupling_xy = torch.stack(
              [
                  torch.einsum("ij,ik->ijk", A_0_i, A_1_i).mean(0)
                  for A_0_i, A_1_i in zip(A_0.split(bs), A_1.split(bs))
              ]
          ).mean(0)
      else:
          coupling_xx = torch.einsum("ij,ik->ijk", A_0, A_0).mean(0)
          coupling_yy = torch.einsum("ij,ik->ijk", A_1, A_1).mean(0)
          coupling_xy = torch.einsum("ij,ik->ijk", A_0, A_1).mean(0)
      sim_xx = (K[idx_0][:, idx_0] * coupling_xx).sum()
      sim_yy = (K[idx_1][:, idx_1] * coupling_yy).sum()
      sim_xy = (K[idx_0][:, idx_1] * coupling_xy).sum()
      stat = sim_xx + sim_yy - 2 * sim_xy
  
      return (
          stat.cpu(),
          coupling_xx.cpu(),
          coupling_yy.cpu(),
          coupling_xy.cpu(),
      )
  
file_path: alibi-detect/alibi_detect/cd/pytorch/context_aware.py
imports: import logging
  import numpy
  import torch
  from typing import Callable, Dict, Optional, Tuple, Union
  from alibi_detect.cd.base import BaseContextMMDDrift
  from alibi_detect.utils.pytorch import get_device
  from alibi_detect.utils.pytorch.kernels import GaussianRBF
  from alibi_detect.utils.warnings import deprecated_alias
  from alibi_detect.utils.frameworks import Framework
  from alibi_detect.cd._domain_clf import _SVCDomainClf
  from alibi_detect.utils._types import TorchDeviceType
  from tqdm import tqdm
label_desc: Use image node_id 3 for calling the ContextMMDDriftTorch obj's underlying member method code with example usage: obj._cmmd(K, L, bools, L_held) and returns: stat,  coupling_xx,  coupling_yy,  coupling_xy