docstr: Train VAE model.
  
  Parameters
  ----------
  X
      Training batch.
  loss_fn
      Loss function used for training.
  optimizer
      Optimizer used for training.
  cov_elbo
      Dictionary with covariance matrix options in case the elbo loss function is used.
      Either use the full covariance matrix inferred from X (dict(cov_full=None)),
      only the variance (dict(cov_diag=None)) or a float representing the same standard deviation
      for each feature (e.g. dict(sim=.05)).
  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.
docstr_tok: ["Train","VAE","model",".","Parameters","--","--","--","--","--","X","Training","batch",".","loss_fn","Loss","function","used","for","training",".","optimizer","Optimizer","used","for","training",".","cov_elbo","Dictionary","with","covariance","matrix","options","in","case","the","elbo","loss","function","is","used",".","Either","use","the","full","covariance","matrix","inferred","from","X","(","dict","(","cov_full=None",")",")",",","only","the","variance","(","dict","(","cov_diag=None",")",")","or","a","float","representing","the","same","standard","deviation","for","each","feature","(","e.g",".","dict","(","sim=.05",")",")",".","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","."]
code: def fit(
      self,
      X: np.ndarray,
      loss_fn: tf.keras.losses = elbo,
      optimizer: OptimizerTF = tf.keras.optimizers.Adam,
      cov_elbo: dict = dict(sim=0.05),
      epochs: int = 20,
      batch_size: int = 64,
      verbose: bool = True,
      log_metric: Tuple[str, "tf.keras.metrics"] = None,
      callbacks: tf.keras.callbacks = None,
  ) -> None:
      """
      Train VAE model.
  
      Parameters
      ----------
      X
          Training batch.
      loss_fn
          Loss function used for training.
      optimizer
          Optimizer used for training.
      cov_elbo
          Dictionary with covariance matrix options in case the elbo loss function is used.
          Either use the full covariance matrix inferred from X (dict(cov_full=None)),
          only the variance (dict(cov_diag=None)) or a float representing the same standard deviation
          for each feature (e.g. dict(sim=.05)).
      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.
      """
      # train arguments
      args = [self.vae, 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,
      }
  
      # initialize covariance matrix if elbo loss fn is used
      use_elbo = loss_fn.__name__ == "elbo"
      cov_elbo_type, cov = [*cov_elbo][0], [*cov_elbo.values()][0]
      if use_elbo and cov_elbo_type in ["cov_full", "cov_diag"]:
          cov = tfp.stats.covariance(X.reshape(X.shape[0], -1))
          if (
              cov_elbo_type == "cov_diag"
          ):  # infer standard deviation from covariance matrix
              cov = tf.math.sqrt(tf.linalg.diag_part(cov))
      if use_elbo:
          kwargs["loss_fn_kwargs"] = {
              cov_elbo_type: tf.dtypes.cast(cov, tf.float32)
          }
  
      # train
      trainer(*args, **kwargs)
  
imports: import logging
  import numpy
  import tensorflow
  import tensorflow_probability
  from typing import Dict, Tuple
  from alibi_detect.models.tensorflow.autoencoder import VAE
  from alibi_detect.models.tensorflow.trainer import trainer
  from alibi_detect.models.tensorflow.losses import elbo
  from alibi_detect.base import BaseDetector, FitMixin, ThresholdMixin, outlier_prediction_dict
  from alibi_detect.utils.tensorflow.prediction import predict_batch
  from alibi_detect.utils._types import OptimizerTF
label_desc: Use image node_id 2 for calling the OutlierVAE obj's underlying member method code with example usage: obj.fit(X, loss_fn, optimizer, cov_elbo, epochs, batch_size, verbose, log_metric, callbacks) without return types