Warn if steps_per_execution truncated, error if not divisor of inferred_steps

Summary:
Reinstates the more informative error ([here in SDK2.4](https://phabricator.sourcevertex.net/diffusion/TENSORFLOW/browse/poplar%252Fr2.4%252Frelease/tensorflow/python/ipu/keras/extensions/data_adapter.py$210)) that lets the user know how they have incorrectly set `steps_per_execution` (if it isn't a divisor of the `inferred_steps` of the dataset). Currently, the error thrown is an Infeed overflow error, which isn't very clear.

Additionally, log a warning if truncation has occurred due to `steps_per_execution` being larger than the `inferred_steps` of the dataset.

TF2.5 Only

Test Plan: Test added to ensure that the error is thrown under the necessary circumstances.

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

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

Subscribers: christiana

Maniphest Tasks: T55559

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

Author: hmellor

tensorflow/compiler/plugin/poplar/tests/distributed_tf2_test.py

@@ -164,16 +164,19 @@ def test_tf2_distributed_popdist_strategy(self):
     strategy = popdist_strategy.PopDistStrategy()
 
     with strategy.scope():
-      dataset = test_dataset(popdist.getNumInstances() * batch_size *
+      dataset = test_dataset(popdist.getNumTotalReplicas() * batch_size *
                              steps_to_run,
                              batch_size=batch_size)
+      dataset = dataset.shard(num_shards=popdist.getNumInstances(),
+                              index=popdist.getInstanceIndex())
+      steps_per_execution = len(dataset) // popdist.getNumLocalReplicas()
       model = simple_model()
       optimizer = tf.keras.optimizers.SGD(learning_rate=0.01)
       loss_fn = tf.keras.losses.MeanSquaredError()
 
       model.compile(optimizer=optimizer,
                     loss=loss_fn,
-                    steps_per_execution=popdist.getNumTotalReplicas())
+                    steps_per_execution=steps_per_execution)
 
       # Build the model separately so we can assert that the biases are
       # broadcasted properly before training.

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

@@ -139,6 +139,15 @@ def __init__(  # pylint: disable=super-init-not-called
   def _validate_data_handler(self):
     super()._validate_data_handler()
 
+    if self.steps_per_execution_value > self.inferred_steps:
+      logging.warn(
+          "`steps_per_execution` has been set to {} but the dataset "
+          "provided{} only contains {} batches. Using {} as "
+          "`steps_per_execution`.".format(
+              self.steps_per_execution_value,
+              " to this replica" if self._replication_factor > 1 else "",
+              self.inferred_steps, self.inferred_steps))
+
     with self._truncate_execution_to_epoch():
       if self.inferred_steps == 0:
         steps_per_replica = math.ceil(
@@ -151,6 +160,11 @@ def _validate_data_handler(self):
                 len(self._dataset), self._replication_factor,
                 steps_per_replica,
                 steps_per_replica * self._replication_factor))
+      elif self.inferred_steps % self._steps_per_execution != 0:
+        raise ValueError(
+            "`steps_per_execution` must be a divisor of the number of batches "
+            "in the dataset provided{}.".format(
+                " to this replica" if self._replication_factor > 1 else ""))
 
   def _infer_steps(self, steps, dataset):
     """Infers steps_per_epoch needed to loop through a dataset."""

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

@@ -17,6 +17,7 @@
 from absl.testing import parameterized
 import numpy as np
 
+from tensorflow.python.ops import variables
 from tensorflow.python.data.experimental.ops import cardinality
 from tensorflow.python.data.ops import dataset_ops
 from tensorflow.python.framework import ops
@@ -275,6 +276,17 @@ def test_dataset_not_big_enough_with_replication(self):
                                                  replication_factor=4)
       del data_handler
 
+  def test_steps_per_execution_not_divisor(self):
+    x = np.ones((8, 1))
+    x = dataset_ops.Dataset.from_tensor_slices(x).batch(2, drop_remainder=True)
+    with self.assertRaisesRegex(
+        ValueError,
+        r"`steps_per_execution` must be a divisor of the number of batches in "
+        r"the dataset provided."):
+      data_handler = data_adapter.IPUDataHandler(
+          x, epochs=1, batch_size=2, steps_per_execution=variables.Variable(3))
+      del data_handler
+
 
 if __name__ == '__main__':
   ops.enable_eager_execution()