docstr: Train Adversarial AE model.
  
  Parameters
  ----------
  X
      Training batch.
  loss_fn
      Loss function used for training.
  w_model
      Weight on model prediction loss term.
  w_recon
      Weight on MSE reconstruction error loss term.
  optimizer
      Optimizer used for training.
  epochs
      Number of training epochs.
  batch_size
      Batch size used for training.
  verbose
      Whether to print training progress.
  log_metric
      Additional metrics whose progress will be displayed if verbose equals True.
  callbacks
      Callbacks used during training.
  preprocess_fn
      Preprocessing function applied to each training batch.
docstr_tok: ["Train","Adversarial","AE","model",".","Parameters","--","--","--","--","--","X","Training","batch",".","loss_fn","Loss","function","used","for","training",".","w_model","Weight","on","model","prediction","loss","term",".","w_recon","Weight","on","MSE","reconstruction","error","loss","term",".","optimizer","Optimizer","used","for","training",".","epochs","Number","of","training","epochs",".","batch_size","Batch","size","used","for","training",".","verbose","Whether","to","print","training","progress",".","log_metric","Additional","metrics","whose","progress","will","be","displayed","if","verbose","equals","True",".","callbacks","Callbacks","used","during","training",".","preprocess_fn","Preprocessing","function","applied","to","each","training","batch","."]
code: def fit(
      self,
      X: np.ndarray,
      loss_fn: tf.keras.losses = loss_adv_ae,
      w_model: float = 1.0,
      w_recon: float = 0.0,
      optimizer: OptimizerTF = tf.keras.optimizers.Adam,
      epochs: int = 20,
      batch_size: int = 128,
      verbose: bool = True,
      log_metric: Tuple[str, "tf.keras.metrics"] = None,
      callbacks: tf.keras.callbacks = None,
      preprocess_fn: Callable = None,
  ) -> None:
      """
      Train Adversarial AE model.
  
      Parameters
      ----------
      X
          Training batch.
      loss_fn
          Loss function used for training.
      w_model
          Weight on model prediction loss term.
      w_recon
          Weight on MSE reconstruction error loss term.
      optimizer
          Optimizer used for training.
      epochs
          Number of training epochs.
      batch_size
          Batch size used for training.
      verbose
          Whether to print training progress.
      log_metric
          Additional metrics whose progress will be displayed if verbose equals True.
      callbacks
          Callbacks used during training.
      preprocess_fn
          Preprocessing function applied to each training batch.
      """
      # train arguments
      args = [self.ae, loss_fn, X]
      optimizer = (
          optimizer() if isinstance(optimizer, type) else optimizer
      )
      kwargs = {
          "optimizer": optimizer,
          "epochs": epochs,
          "batch_size": batch_size,
          "verbose": verbose,
          "log_metric": log_metric,
          "callbacks": callbacks,
          "preprocess_fn": preprocess_fn,
          "loss_fn_kwargs": {
              "model": self.model,
              "model_hl": self.model_hl,
              "w_model": w_model,
              "w_recon": w_recon,
              "w_model_hl": self.w_model_hl,
              "temperature": self.temperature,
          },
      }
  
      # train
      trainer(*args, **kwargs)
  
imports: import logging
  from typing import Callable, Dict, List, Tuple, Union, cast
  import numpy
  import tensorflow
  from alibi_detect.base import BaseDetector, FitMixin, ThresholdMixin, adversarial_correction_dict, adversarial_prediction_dict
  from alibi_detect.models.tensorflow.autoencoder import AE
  from alibi_detect.models.tensorflow.losses import loss_adv_ae
  from alibi_detect.models.tensorflow.trainer import trainer
  from alibi_detect.utils.tensorflow.prediction import predict_batch
  from alibi_detect.utils._types import OptimizerTF
  from tensorflow.keras.layers import Dense, Flatten
  from tensorflow.keras.losses import kld
  from tensorflow.keras.models import Model
label_desc: Use image node_id 2 for calling the AdversarialAE obj's underlying member method code with example usage: obj.fit(X, loss_fn, w_model, w_recon, optimizer, epochs, batch_size, verbose, log_metric, callbacks, preprocess_fn) without return types