code: def eval_step(state, data):
      (
          trainable_variables,
          non_trainable_variables,
          metric_variables,
      ) = state
      x, y = data
      y_pred, non_trainable_variables = model.stateless_call(
          trainable_variables, non_trainable_variables, x
      )
      loss = loss_fn(y, y_pred)
      metric_variables = val_acc_metric.stateless_update_state(
          metric_variables, y, y_pred
      )
      return loss, (
          trainable_variables,
          non_trainable_variables,
          metric_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 7 for calling a global function with example usage: eval_step(state, data) and returns: loss,  trainable_variables,  non_trainable_variables,  metric_variables