code: def mnist_dataset(batch_size: int) -> tf.data.Dataset:
      (
          x_train,
          y_train,
      ), _ = tf.keras.datasets.fashion_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, seed=1234)
          .batch(batch_size)
      )
      return train_dataset
  
file_path: ray/release/air_tests/air_benchmarks/workloads/tensorflow_benchmark.py
imports: import json
  import os
  import time
  from pathlib import Path
  import click
  import numpy
  import tensorflow
  from typing import List, Tuple
label_desc: Use image node_id 1 for calling a global function with example usage: mnist_dataset(batch_size) and returns: train_dataset