code: def loss_adv_ae(
      x_true: tf.Tensor,
      x_pred: tf.Tensor,
      model: tf.keras.Model = None,
      model_hl: list = None,
      w_model: float = 1.0,
      w_recon: float = 0.0,
      w_model_hl: list = None,
      temperature: float = 1.0,
  ) -> tf.Tensor:
      """
      Loss function used for AdversarialAE.
  
      Parameters
      ----------
      x_true
          Batch of instances.
      x_pred
          Batch of reconstructed instances by the autoencoder.
      model
          A trained tf.keras model with frozen layers (layers.trainable = False).
      model_hl
          List with tf.keras models used to extract feature maps and make predictions on hidden layers.
      w_model
          Weight on model prediction loss term.
      w_recon
          Weight on MSE reconstruction error loss term.
      w_model_hl
          Weights assigned to the loss of each model in model_hl.
      temperature
          Temperature used for model prediction scaling.
          Temperature <1 sharpens the prediction probability distribution.
  
      Returns
      -------
      Loss value.
      """
      y_true = model(x_true)
      y_pred = model(x_pred)
  
      # apply temperature scaling
      if temperature != 1.0:
          y_true = y_true ** (1 / temperature)
          y_true = y_true / tf.reshape(
              tf.reduce_sum(y_true, axis=-1), (-1, 1)
          )
  
      # compute K-L divergence loss
      loss_kld = kld(y_true, y_pred)
      std_kld = tf.math.reduce_std(loss_kld)
      loss = tf.reduce_mean(loss_kld)
  
      # add loss from optional K-L divergences extracted from hidden layers
      if isinstance(model_hl, list):
          if w_model_hl is None:
              w_model_hl = list(tf.ones(len(model_hl)))
          for m, w in zip(model_hl, w_model_hl):
              h_true = m(x_true)
              h_pred = m(x_pred)
              loss_kld_hl = tf.reduce_mean(kld(h_true, h_pred))
              loss += tf.constant(w) * loss_kld_hl
      loss *= w_model
  
      # add optional reconstruction loss
      if w_recon > 0.0:
          loss_recon = (x_true - x_pred) ** 2
          std_recon = tf.math.reduce_std(loss_recon)
          w_scale = std_kld / (std_recon + 1e-10)
          loss_recon = w_recon * w_scale * tf.reduce_mean(loss_recon)
          loss += loss_recon
          return loss
      else:
          return loss
  
file_path: alibi-detect/alibi_detect/models/tensorflow/losses.py
imports: from typing import Optional
  import tensorflow
  from tensorflow.keras.layers import Flatten
  from tensorflow.keras.losses import kld, categorical_crossentropy
  import tensorflow_probability
  from alibi_detect.models.tensorflow.gmm import gmm_params, gmm_energy
label_desc: Use image node_id 4 for calling a global function with example usage: loss_adv_ae(x_true, x_pred, model, model_hl, w_model, w_recon, w_model_hl, temperature) and returns: loss, loss