code: def get_datasets():
      # Load the data and split it between train and test sets
      (x_train, y_train), (
          x_test,
          y_test,
      ) = keras.datasets.mnist.load_data()
  
      # Scale images to the [0, 1] range
      x_train = x_train.astype("float32")
      x_test = x_test.astype("float32")
      # Make sure images have shape (28, 28, 1)
      x_train = np.expand_dims(x_train, -1)
      x_test = np.expand_dims(x_test, -1)
      print("x_train shape:", x_train.shape)
      print(x_train.shape[0], "train samples")
      print(x_test.shape[0], "test samples")
  
      # Create TF Datasets
      train_data = tf.data.Dataset.from_tensor_slices(
          (x_train, y_train)
      )
      eval_data = tf.data.Dataset.from_tensor_slices((x_test, y_test))
      return train_data, eval_data
  
file_path: keras/guides/distributed_training_with_jax.py
imports: import os
  import jax
  import numpy
  import tensorflow
  import keras
  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 2 for calling a global function with example usage: get_datasets() and returns: train_data,  eval_data