code: def mnist_dataset(batch_size: int) -> tf.data.Dataset:
      with FileLock(os.path.expanduser("~/.mnist_lock")):
          (x_train, y_train), _ = tf.keras.datasets.mnist.load_data()
      # The `x` arrays are in uint8 and have values in the [0, 255] range.
      # You need to convert them to float32 with values in the [0, 1] range.
      x_train = x_train / np.float32(255)
      y_train = y_train.astype(np.int64)
      train_dataset = (
          tf.data.Dataset.from_tensor_slices((x_train, y_train))
          .shuffle(60000)
          .repeat()
          .batch(batch_size)
      )
      return train_dataset
  
file_path: ray/python/ray/train/examples/tf/tensorflow_mnist_example.py
imports: import argparse
  import json
  import os
  import numpy
  import tensorflow
  from filelock import FileLock
  from ray.air.integrations.keras import ReportCheckpointCallback
  from ray.train import Result, RunConfig, ScalingConfig
  from ray.train.tensorflow import TensorflowTrainer
label_desc: Use image node_id 1 for calling a global function with example usage: mnist_dataset(batch_size) and returns: train_dataset