code: def train_func(config: dict):
      per_worker_batch_size = config.get("batch_size", 64)
      epochs = config.get("epochs", 3)
  
      dataset_shard = train.get_dataset_shard("train")
  
      strategy = tf.distribute.MultiWorkerMirroredStrategy()
  
      with strategy.scope():
          # Model building/compiling need to be within `strategy.scope()`.
          multi_worker_model = build_autoencoder_model()
          learning_rate = config.get("lr", 0.001)
          multi_worker_model.compile(
              loss=tf.keras.losses.BinaryCrossentropy(),
              optimizer=tf.keras.optimizers.Adam(
                  learning_rate=learning_rate
              ),
              metrics=[
                  "binary_crossentropy",
              ],
          )
  
      def to_tf_dataset(dataset, batch_size):
          def to_tensor_iterator():
              for batch in dataset.iter_tf_batches(
                  batch_size=batch_size, dtypes=tf.float32
              ):
                  yield batch["image"], batch["label"]
  
          output_signature = (
              tf.TensorSpec(shape=(None, 784), dtype=tf.float32),
              tf.TensorSpec(shape=(None, 784), dtype=tf.float32),
          )
          tf_dataset = tf.data.Dataset.from_generator(
              to_tensor_iterator, output_signature=output_signature
          )
          return prepare_dataset_shard(tf_dataset)
  
      results = []
      for epoch in range(epochs):
          tf_dataset = to_tf_dataset(
              dataset=dataset_shard,
              batch_size=per_worker_batch_size,
          )
          history = multi_worker_model.fit(
              tf_dataset, callbacks=[ReportCheckpointCallback()]
          )
          results.append(history.history)
      return results
  
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 3 for calling a global function with example usage: train_func(config) and returns: results, prepare_dataset_shard