docstr: Train semantic and background generative models.
  
  Parameters
  ----------
  X
      Training batch.
  mutate_fn
      Mutation function used to generate the background dataset.
  mutate_fn_kwargs
      Kwargs for the mutation function used to generate the background dataset.
      Default values set for an image dataset.
  mutate_batch_size
      Batch size used to generate the mutations for the background dataset.
  loss_fn
      Loss function used for training.
  loss_fn_kwargs
      Kwargs for loss function.
  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.
docstr_tok: ["Train","semantic","and","background","generative","models",".","Parameters","--","--","--","--","--","X","Training","batch",".","mutate_fn","Mutation","function","used","to","generate","the","background","dataset",".","mutate_fn_kwargs","Kwargs","for","the","mutation","function","used","to","generate","the","background","dataset",".","Default","values","set","for","an","image","dataset",".","mutate_batch_size","Batch","size","used","to","generate","the","mutations","for","the","background","dataset",".","loss_fn","Loss","function","used","for","training",".","loss_fn_kwargs","Kwargs","for","loss","function",".","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","."]
code: def fit(
      self,
      X: np.ndarray,
      mutate_fn: Callable = mutate_categorical,
      mutate_fn_kwargs: dict = {
          "rate": 0.2,
          "seed": 0,
          "feature_range": (0, 255),
      },
      mutate_batch_size: int = int(1e10),
      loss_fn: tf.keras.losses = None,
      loss_fn_kwargs: dict = None,
      optimizer: OptimizerTF = tf.keras.optimizers.Adam,
      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 semantic and background generative models.
  
      Parameters
      ----------
      X
          Training batch.
      mutate_fn
          Mutation function used to generate the background dataset.
      mutate_fn_kwargs
          Kwargs for the mutation function used to generate the background dataset.
          Default values set for an image dataset.
      mutate_batch_size
          Batch size used to generate the mutations for the background dataset.
      loss_fn
          Loss function used for training.
      loss_fn_kwargs
          Kwargs for loss function.
      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.
      """
      input_shape = X.shape[1:]
      optimizer = (
          optimizer() if isinstance(optimizer, type) else optimizer
      )
      # Separate into two separate optimizers, one for semantic model and one for background model
      optimizer_s = optimizer
      optimizer_b = optimizer.__class__.from_config(
          optimizer.get_config()
      )
  
      # training arguments
      kwargs = {
          "epochs": epochs,
          "batch_size": batch_size,
          "verbose": verbose,
          "callbacks": callbacks,
      }
  
      # create background data
      mutate_fn = partial(mutate_fn, **mutate_fn_kwargs)
      X_back = predict_batch(
          X, mutate_fn, batch_size=mutate_batch_size, dtype=X.dtype
      )
  
      # prepare sequential data
      if self.sequential and not self.has_log_prob:
          y, y_back = X[:, 1:], X_back[:, 1:]  # type: ignore
          X, X_back = X[:, :-1], X_back[:, :-1]  # type: ignore
      else:
          y, y_back = None, None
  
      # check if model needs to be built
      use_build = (
          True
          if self.has_log_prob
          and not isinstance(self.dist_s, tf.keras.Model)
          else False
      )
  
      if use_build:
          # build and train semantic model
          self.model_s = build_model(self.dist_s, input_shape)[0]
          self.model_s.compile(optimizer=optimizer_s)
          self.model_s.fit(X, **kwargs)
          # build and train background model
          self.model_b = build_model(self.dist_b, input_shape)[0]
          self.model_b.compile(optimizer=optimizer_b)
          self.model_b.fit(X_back, **kwargs)
      else:
          # update training arguments
          kwargs.update(
              {
                  "loss_fn_kwargs": loss_fn_kwargs,
                  "log_metric": log_metric,
              }
          )
  
          # train semantic model
          args = [self.dist_s, loss_fn, X]
          kwargs.update({"y_train": y, "optimizer": optimizer_s})
          trainer(*args, **kwargs)
  
          # train background model
          args = [self.dist_b, loss_fn, X_back]
          kwargs.update({"y_train": y_back, "optimizer": optimizer_b})
          trainer(*args, **kwargs)
  
imports: from functools import partial
  import logging
  import numpy
  import tensorflow
  from tensorflow.keras.layers import Input
  from tensorflow.keras.models import Model
  from tensorflow_probability.python.distributions.distribution import Distribution
  from typing import Callable, Dict, Tuple, Union
  from alibi_detect.models.tensorflow.pixelcnn import PixelCNN
  from alibi_detect.models.tensorflow.trainer import trainer
  from alibi_detect.base import BaseDetector, FitMixin, ThresholdMixin, outlier_prediction_dict
  from alibi_detect.utils.tensorflow.prediction import predict_batch
  from alibi_detect.utils.tensorflow.perturbation import mutate_categorical
  from alibi_detect.utils._types import OptimizerTF
label_desc: Use image node_id 2 for calling the LLR obj's underlying member method code with example usage: obj.fit(X, mutate_fn, mutate_fn_kwargs, mutate_batch_size, loss_fn, loss_fn_kwargs, optimizer, epochs, batch_size, verbose, log_metric, callbacks) without return types