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: tf.Tensor,
      L: tf.Tensor,
      bools: tf.Tensor,
      L_held: tf.Tensor = None,
  ) -> Tuple[tf.Tensor, tf.Tensor, tf.Tensor, tf.Tensor]:
      """
      Private method to compute the MMD-ADiTT test statistic.
      """
      # Get ref/test indices
      idx_0, idx_1 = np.where(bools == 0)[0], np.where(bools == 1)[0]
      n_ref, n_test = len(idx_0), len(idx_1)
  
      # Form kernel matrices
      L_0, L_1 = tf.gather(
          tf.gather(L, idx_0), idx_0, axis=1
      ), tf.gather(tf.gather(L, idx_1), idx_1, axis=1)
      K_0, K_1 = tf.gather(
          tf.gather(K, idx_0), idx_0, axis=1
      ), tf.gather(tf.gather(K, idx_1), idx_1, axis=1)
      # Avoid using tf.gather_nd since this would require [n_fef, n_ref, 2] and [n_test, n_test, 2] idx tensors
  
      # Initialise regularisation parameters
      # Implemented only for first _cmmd call which corresponds to original window assignment
      if self.lams is None:
          possible_lams = tf.convert_to_tensor(
              [2 ** (-i) for i in range(20)], dtype=tf.float64
          )
          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 = tf.linalg.inv(
          L_0 + n_ref * self.lams[0] * tf.eye(int(n_ref))
      )
      L_1_inv = tf.linalg.inv(
          L_1 + n_test * self.lams[1] * tf.eye(int(n_test))
      )
      A_0 = tf.gather(L_held, idx_0, axis=1) @ L_0_inv
      A_1 = tf.gather(L_held, idx_1, axis=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 = tf.reduce_mean(
              tf.stack(
                  [
                      tf.reduce_mean(
                          tf.einsum("ij,ik->ijk", A_0_i, A_0_i), axis=0
                      )
                      for A_0_i in tf.split(
                          A_0, _split_chunks(len(A_0), bs)
                      )
                  ]
              ),
              axis=0,
          )
          coupling_yy = tf.reduce_mean(
              tf.stack(
                  [
                      tf.reduce_mean(
                          tf.einsum("ij,ik->ijk", A_1_i, A_1_i), axis=0
                      )
                      for A_1_i in tf.split(
                          A_1, _split_chunks(len(A_1), bs)
                      )
                  ]
              ),
              axis=0,
          )
          coupling_xy = tf.reduce_mean(
              tf.stack(
                  [
                      tf.reduce_mean(
                          tf.einsum("ij,ik->ijk", A_0_i, A_1_i), axis=0
                      )
                      for A_0_i, A_1_i in zip(
                          tf.split(A_0, _split_chunks(len(A_0), bs)),
                          tf.split(A_1, _split_chunks(len(A_1), bs)),
                      )
                  ]
              ),
              axis=0,
          )
      else:
          coupling_xx = tf.reduce_mean(
              tf.einsum("ij,ik->ijk", A_0, A_0), axis=0
          )
          coupling_yy = tf.reduce_mean(
              tf.einsum("ij,ik->ijk", A_1, A_1), axis=0
          )
          coupling_xy = tf.reduce_mean(
              tf.einsum("ij,ik->ijk", A_0, A_1), axis=0
          )
      sim_xx = tf.reduce_sum(
          tf.gather(tf.gather(K, idx_0), idx_0, axis=1) * coupling_xx
      )
      sim_yy = tf.reduce_sum(
          tf.gather(tf.gather(K, idx_1), idx_1, axis=1) * coupling_yy
      )
      sim_xy = tf.reduce_sum(
          tf.gather(tf.gather(K, idx_0), idx_1, axis=1) * coupling_xy
      )
      stat = sim_xx + sim_yy - 2 * sim_xy
  
      return stat, coupling_xx, coupling_yy, coupling_xy
  
file_path: alibi-detect/alibi_detect/cd/tensorflow/context_aware.py
imports: import logging
  import numpy
  import tensorflow
  import tensorflow_probability
  from typing import Callable, Dict, Optional, Tuple, Union, List
  from alibi_detect.cd.base import BaseContextMMDDrift
  from alibi_detect.utils.tensorflow.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 tqdm import tqdm
label_desc: Use image node_id 3 for calling the ContextMMDDriftTF 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