Handle Checkpoints in strategy

Summary:
This is due to upstream change 2c9ffb5

Ref T53780

TF2.5 Only

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

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

Maniphest Tasks: T53780

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

Author: georgepaw

tensorflow/python/ipu/ipu_multi_worker_strategy.py

@@ -33,6 +33,7 @@
 from tensorflow.python.ops import control_flow_util
 from tensorflow.python.ops import control_flow_ops
 from tensorflow.python.ops import variable_scope
+from tensorflow.python.training.tracking import base as trackable
 from tensorflow.python.util import tf_contextlib
 from tensorflow.python.util import deprecation
 
@@ -335,6 +336,8 @@ def initial_value_fn():  # pylint: disable=g-missing-docstring
         initial_value = kwargs["initial_value"]
         if callable(initial_value):
           initial_value = initial_value()
+        if isinstance(initial_value, trackable.CheckpointInitialValue):
+          initial_value = initial_value.wrapped_value
         assert not callable(initial_value)
         initial_value = ops.convert_to_tensor(initial_value,
                                               dtype=kwargs.get("dtype", None))

tensorflow/python/ipu/ipu_strategy.py

@@ -36,6 +36,7 @@
 from tensorflow.python.keras.engine import sequential
 from tensorflow.python.ipu.ops import cross_replica_ops
 from tensorflow.python.ops import control_flow_ops
+from tensorflow.python.training.tracking import base as trackable
 from tensorflow.python.util import nest
 from tensorflow.python.ipu.keras.extensions import functional_extensions
 from tensorflow.python.ipu.keras.extensions import sequential_extensions
@@ -126,6 +127,8 @@ def initial_value_fn():
       initial_value = kwargs["initial_value"]
       if callable(initial_value):
         initial_value = initial_value()
+      if isinstance(initial_value, trackable.CheckpointInitialValue):
+        initial_value = initial_value.wrapped_value
       assert not callable(initial_value)
       return ops.convert_to_tensor(initial_value,
                                    dtype=kwargs.get("dtype", None))