Don't try and get length for inifinite datasets

Summary:
Fix T57598
TF2.5 Only

Test Plan: Added a test

Reviewers: christiana, jackh, #tensorflow, #framework_ip_review_-_any_oss_or_third-party_code_use_has_been_approved

Reviewed By: christiana, #tensorflow, #framework_ip_review_-_any_oss_or_third-party_code_use_has_been_approved

Subscribers: harrym

Maniphest Tasks: T57598

Differential Revision: https://phabricator.sourcevertex.net/D62644

Author: georgepaw

tensorflow/python/ipu/keras/extensions/data_adapter.py

@@ -174,12 +174,14 @@ def _infer_steps(self, steps, dataset):
           "Could not infer the size of the data. You must specify the number "
           "of steps to run.")
     if steps % self._replication_factor:
-      logging.warn(
-          "Dataset of length {} is being evenly distributed between {} "
-          "replicas. The remaining {} batch{} will be dropped.".format(
-              len(dataset), self._replication_factor,
-              steps % self._replication_factor,
-              "es" if steps % self._replication_factor > 1 else ""))
+      size = cardinality.cardinality(dataset)
+      if size >= 0:
+        logging.warn(
+            "Dataset of length {} is being evenly distributed between {} "
+            "replicas. The remaining {} batch{} will be dropped.".format(
+                len(dataset), self._replication_factor,
+                steps % self._replication_factor,
+                "es" if steps % self._replication_factor > 1 else ""))
 
     return int(steps // self._replication_factor)
 

tensorflow/python/ipu/tests/keras/extensions/data_adapter_test.py

@@ -287,6 +287,15 @@ def test_steps_per_execution_not_divisor(self):
           x, epochs=1, batch_size=2, steps_per_execution=variables.Variable(3))
       del data_handler
 
+  def test_dataset_drop_batch_with_replication(self):
+    x = np.ones((6, 1))
+    x = dataset_ops.Dataset.from_tensor_slices(x).repeat()
+    data_handler = data_adapter.IPUDataHandler(x,
+                                               steps_per_epoch=7,
+                                               replication_factor=4)
+    del data_handler
+    # No exception raised.
+
 
 if __name__ == '__main__':
   ops.enable_eager_execution()