code: def train_step(state, data):
      (
          trainable_variables,
          non_trainable_variables,
          optimizer_variables,
      ) = state
      x, y = data
      (loss, non_trainable_variables), grads = grad_fn(
          trainable_variables, non_trainable_variables, x, y
      )
      (
          trainable_variables,
          optimizer_variables,
      ) = optimizer.stateless_apply(
          optimizer_variables, grads, trainable_variables
      )
      # Return updated state
      return loss, (
          trainable_variables,
          non_trainable_variables,
          optimizer_variables,
      )
  
file_path: keras/guides/writing_a_custom_training_loop_in_jax.py
imports: import os
  import jax
  import tensorflow
  import keras
  import numpy
label_desc: Use image node_id 4 for calling a global function with example usage: train_step(state, data) and returns: loss,  trainable_variables,  non_trainable_variables,  optimizer_variables