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