code: def build_train_step_fn(
      model, optimizer, loss_fn, metrics, attack_fn
  ):
      """Builds a compiled tf.function for training the model for one step.
  
      Args:
        model: A Keras Model object which makes inference on input.
        optimizer: A Keras Optimizer object to be used for training.
        loss_fn: A callable representing the loss objective for training.
        metrics: A list of Keras Metric objects to be updated during training.
        attack_fn: A callable built by build_attack_fn.
  
      Returns:
        A compiled tf.function which trains the model for one step.
      """
  
      @tf.function
      def train_step(x, y):
          adv_x = attack_fn(x, y)
          with tf.GradientTape() as tape:
              logits = model(adv_x, training=True)
              loss_value = loss_fn(y, logits) + sum(model.losses)
          trainable_vars = model.trainable_variables
          gradients = tape.gradient(loss_value, trainable_vars)
          optimizer.apply_gradients(zip(gradients, trainable_vars))
          for metric in metrics:
              metric.update_state(y, logits)
          return loss_value
  
      return train_step
  
file_path: neural-structured-learning/research/multi_representation_adversary/multi_representation_adversary/helper.py
imports: from absl import logging
  from multi_representation_adversary import attacks
  from multi_representation_adversary import transforms
  import tensorflow.compat.v2
label_desc: Use image node_id 2 for calling a global function with example usage: build_train_step_fn(model, optimizer, loss_fn, metrics, attack_fn) and returns: train_step, loss_value