class_docstr: Abstract custom Keras model.
  
   This model overwrites the following methods:
  - train_step
  - test_step
  - predict_step
  - save
  - load
  Cannot be used as tf.keras.Model.
class_docstr_tok: ["Abstract","custom","Keras","model",".","This","model","overwrites","the","following","methods",":","-","train_step","-","test_step","-","predict_step","-","save","-","load","Can","not","be","used","as","tf.keras.Model","."]
docstr: Performs a train step using the given batch.
  
  Args:
      batch_in: The batch input.
  
  Returns:
      Training metrics.
docstr_tok: ["Performs","a","train","step","using","the","given","batch",".","Args",":","batch_in",":","The","batch","input",".","Returns",":","Training","metrics","."]
code: def train_step(
      self,
      batch_in: Union[Tuple[tf.Tensor, ...], Tuple[np.ndarray, ...]],
  ) -> Dict[Text, float]:
      """Performs a train step using the given batch.
  
      Args:
          batch_in: The batch input.
  
      Returns:
          Training metrics.
      """
      self._training = True
  
      # calculate supervision and regularization losses separately
      with tf.GradientTape(persistent=True) as tape:
          prediction_loss = self.batch_loss(batch_in)
          regularization_loss = tf.math.add_n(self.losses)
          total_loss = prediction_loss + regularization_loss
  
      self.total_loss.update_state(total_loss)
  
      # calculate the gradients that come from supervision signal
      prediction_gradients = tape.gradient(
          prediction_loss, self.trainable_variables
      )
      # calculate the gradients that come from regularization
      regularization_gradients = tape.gradient(
          regularization_loss, self.trainable_variables
      )
      # delete gradient tape manually
      # since it was created with `persistent=True` option
      del tape
  
      gradients = []
      for pred_grad, reg_grad in zip(
          prediction_gradients, regularization_gradients
      ):
          if pred_grad is not None and reg_grad is not None:
              # remove regularization gradient for variables
              # that don't have prediction gradient
              gradients.append(
                  pred_grad
                  + tf.where(
                      pred_grad > 0, reg_grad, tf.zeros_like(reg_grad)
                  )
              )
          else:
              gradients.append(pred_grad)
  
      self.optimizer.apply_gradients(
          zip(gradients, self.trainable_variables)
      )
  
      self._training = None
  
      return self._get_metric_results()
  
imports: import time
  import random
  import tensorflow
  import numpy
  import logging
  import os
  from collections import defaultdict
  from typing import List, Text, Dict, Tuple, Union, Optional, Any, TYPE_CHECKING
  from keras.utils import tf_utils
  from keras import Model
  from rasa.shared.constants import DIAGNOSTIC_DATA
  from rasa.utils.tensorflow.constants import LABEL, IDS, INTENT_CLASSIFICATION, SENTENCE, SEQUENCE_LENGTH, RANDOM_SEED, EMBEDDING_DIMENSION, REGULARIZATION_CONSTANT, SIMILARITY_TYPE, CONNECTION_DENSITY, NUM_NEG, LOSS_TYPE, MAX_POS_SIM, MAX_NEG_SIM, USE_MAX_NEG_SIM, NEGATIVE_MARGIN_SCALE, SCALE_LOSS, LEARNING_RATE, CONSTRAIN_SIMILARITIES, MODEL_CONFIDENCE, RUN_EAGERLY
  from rasa.utils.tensorflow.model_data import RasaModelData, FeatureSignature, FeatureArray
  import rasa.utils.train_utils
  from rasa.utils.tensorflow import layers
  from rasa.utils.tensorflow import rasa_layers
  from rasa.utils.tensorflow.data_generator import RasaDataGenerator, RasaBatchDataGenerator
  from rasa.shared.nlu.constants import TEXT
  from rasa.shared.exceptions import RasaException
  from rasa.utils.tensorflow.types import BatchData, MaybeNestedBatchData
label_desc: Use image node_id 6 for calling the RasaModel obj's underlying member method code with example usage: obj.train_step(batch_in) and returns: self