Add a require_static_shapes argument to DistributionStrategy class. This allows us to identify if we need to set the drop_remainder option when creating Dataset objects.

PiperOrigin-RevId: 215633097

1 files changed, 7 insertions, 0 deletions

diff --git a/tensorflow/python/training/distribute.py b/tensorflow/python/training/distribute.py
index a92a1bdee7..b3f3c29b2f 100644
--- a/tensorflow/python/training/distribute.py
+++ b/tensorflow/python/training/distribute.py
@@ -436,6 +436,9 @@ class DistributionStrategy(object):
 
   def __init__(self):
     self._default_device = None
+    # This property is used to determine if we should set drop_remainder=True
+    # when creating Datasets from numpy array inputs.
+    self._require_static_shapes = False
 
   def scope(self):
     """Returns a context manager selecting this DistributionStrategy as current.
@@ -899,6 +902,10 @@ class DistributionStrategy(object):
     raise NotImplementedError("must be implemented in descendants")
 
   @property
+  def require_static_shapes(self):
+    return self._require_static_shapes
+
+  @property
   def num_towers(self):
     """Returns number of towers, for purposes of averaging across towers."""
     raise NotImplementedError("must be implemented in descendants")