code: def train_step(train_state, x, y):
      (loss_value, non_trainable_variables), grads = compute_gradients(
          train_state.trainable_variables,
          train_state.non_trainable_variables,
          x,
          y,
      )
  
      (
          trainable_variables,
          optimizer_variables,
      ) = optimizer.stateless_apply(
          train_state.optimizer_variables,
          grads,
          train_state.trainable_variables,
      )
  
      return loss_value, TrainingState(
          trainable_variables,
          non_trainable_variables,
          optimizer_variables,
      )
  
file_path: keras/examples/demo_jax_distributed.py
imports: import os, pprint, collections
  import jax
  import jax.numpy
  import tensorflow
  import keras
  import numpy
  from tqdm import tqdm
  from jax.experimental import mesh_utils
  from jax.sharding import Mesh
  from jax.sharding import NamedSharding
  from jax.sharding import PartitionSpec
label_desc: Use image node_id 4 for calling a global function with example usage: train_step(train_state, x, y) and returns: loss_value,  TrainingState