code: def test_orca_tf_dataset(self):
      sc = OrcaContext.get_spark_context()
      rdd = sc.parallelize(range(5))
      shard = SparkXShards(rdd)
  
      def generate_dataset(df, a, b, size):
          items = np.array(
              [[i / size] for i in range(size)], dtype=np.float32
          )
          labels = np.array(
              [a * x[0] + b for x in items], dtype=np.float32
          )
          return {"item": items, "label": labels}
  
      train_shard = shard.transform_shard(generate_dataset, 5, 10, 1000)
      test_shard = shard.transform_shard(generate_dataset, 5, 10, 100)
  
      train_dataset = Dataset.from_tensor_slices(train_shard)
      test_dataset = Dataset.from_tensor_slices(test_shard)
      train_step = math.ceil(5 * 1000 / 32)
      test_step = math.ceil(5 * 100 / 32)
  
      orca_estimator = get_estimator(
          workers_per_node=2, model_fn=model_creator_for_orca_dataset
      )
  
      start_eval_stats = orca_estimator.evaluate(
          data=test_dataset, num_steps=test_step, batch_size=32
      )
  
      train_stats = orca_estimator.fit(
          data=train_dataset,
          epochs=2,
          batch_size=32,
          steps_per_epoch=train_step,
      )
      print(train_stats)
  
      end_eval_stats = orca_estimator.evaluate(
          data=test_dataset, num_steps=test_step, batch_size=32
      )
  
      assert isinstance(train_stats, dict), "fit should return a dict"
      assert isinstance(
          end_eval_stats, dict
      ), "evaluate should return a dict"
      assert orca_estimator.get_model(
          sample_input={"item": np.array([[1]], dtype=np.float32)}
      )
  
      dloss = (
          end_eval_stats["validation_loss"]
          - start_eval_stats["validation_loss"]
      )
      dmse = (
          end_eval_stats["validation_mean_squared_error"]
          - start_eval_stats["validation_mean_squared_error"]
      )
      print(f"dLoss: {dloss}, dMSE: {dmse}")
      assert (
          dloss < 0 and dmse < 0
      ), "training sanity check failed. loss increased!"
  
      pred_shards = orca_estimator.predict(test_dataset)
      pred1 = pred_shards.collect()
  
      path = "/tmp/model_test_orca_dataset.ckpt"
      try:
          orca_estimator.save_checkpoint(path)
          orca_estimator.shutdown()
          est = get_estimator(model_fn=model_creator_for_orca_dataset)
          with self.assertRaises(Exception) as context:
              est.load_checkpoint(path)
          self.assertTrue(
              "Failed to set model weights, please provide real tensor data"
              in str(context.exception)
          )
          est.load_checkpoint(
              path,
              sample_input={"item": np.array([[1]], dtype=np.float32)},
          )
          result_shards = est.predict(test_dataset)
          pred2 = result_shards.collect()
      finally:
          os.remove(path)
  
      assert np.allclose(pred1[0]["prediction"], pred2[0]["prediction"])
  
imports: import math
  import os
  import shutil
  import tempfile
  import pytest
  from unittest import TestCase
  import tensorflow
  import numpy
  from bigdl.orca.data.pandas import read_csv
  from bigdl.orca.data import SparkXShards, XShards
  from bigdl.orca.data.tf.data import Dataset
  from bigdl.orca.learn.tf2 import Estimator
  from bigdl.orca.ray import OrcaRayContext
  from bigdl.orca import OrcaContext
label_desc: Use image node_id 1 for calling the TestTF2EstimatorRayBackend obj's underlying member method code with example usage: obj.test_orca_tf_dataset() and returns: dict