Add the unit test for a fixed big.

PiperOrigin-RevId: 341515086
Change-Id: Ia6ba3e17a44d8d619748006f414bb217bc703ba7

tensorflow/python/distribute/coordinator/cluster_coordinator_test.py

@@ -658,6 +658,26 @@ class ClusterCoordinatorTest(TestCaseWithErrorReportingThread):
         'only accepts a `tf.function` or a concrete function.'):
       self.coordinator.schedule(func, args=(1,))
 
+  def testDatasetPartiallyCreatedOnCoordinator(self):
+    dataset = dataset_ops.DatasetV2.range(1, 10)
+
+    @def_function.function
+    def input_fn():
+      return dataset.shuffle(9)
+
+    @def_function.function
+    def worker_fn(iterator):
+      x = next(iterator)
+      return x
+
+    per_worker_dataset = self.coordinator.create_per_worker_dataset(input_fn)
+    self.coordinator.schedule(worker_fn, args=(iter(per_worker_dataset),))
+
+    with self.assertRaisesRegexp(
+        coordinator_lib.InputError,
+        'error message is Failed copying input tensor from'):
+      self.coordinator.join()
+
 
 class LimitedClosureQueueSizeBasicTest(ClusterCoordinatorTest):
   """Test basic functionality works with explicit maximum closure queue size.