code: def get_dataset(split_type="train"):
      def dataset_factory():
          return tfds.load(
              "mnist", split=[split_type], as_supervised=True
          )[0].take(128)
  
      dataset = ray.data.read_datasource(
          SimpleTensorFlowDatasource(), dataset_factory=dataset_factory
      )
  
      def normalize_images(x):
          x = np.float32(x.numpy()) / 255.0
          x = np.reshape(x, (-1,))
          return x
  
      def preprocess_dataset(batch):
          return [
              (normalize_images(image), normalize_images(image))
              for image, _ in batch
          ]
  
      dataset = dataset.map_batches(preprocess_dataset)
  
      def convert_batch_to_pandas(batch):
          images = [TensorArray(image) for image, _ in batch]
          # because we did autoencoder here
          df = pd.DataFrame({"image": images, "label": images})
          return df
  
      dataset = dataset.map_batches(convert_batch_to_pandas)
      return dataset
  
file_path: ray/python/ray/train/examples/tf/tensorflow_autoencoder_example.py
imports: import argparse
  import numpy
  import pandas
  import tensorflow
  import tensorflow_datasets
  import ray
  from ray import train
  from ray.air.integrations.keras import ReportCheckpointCallback
  from ray.data.datasource import SimpleTensorFlowDatasource
  from ray.data.extensions import TensorArray
  from ray.train import Result
  from ray.train.tensorflow import TensorflowTrainer, prepare_dataset_shard
label_desc: Use image node_id 1 for calling a global function with example usage: get_dataset(split_type) and returns: dataset, tfds, x, unknown, df