Add basic serialization support to DistributedVariable (by using the underlying primary variable's serialization). Also, throw an exception when trying to de-serialize as we haven't implemented that yet.

guptapriya · tensorflower-gardener · commit 28dec7f4669e · 2018-03-29T19:27:22.000-07:00
PiperOrigin-RevId: 191022884
diff --git a/tensorflow/contrib/distribute/python/values.py b/tensorflow/contrib/distribute/python/values.py
@@ -216,6 +216,9 @@ def shape(self):
   def get_shape(self):
     return self._primary_var.get_shape()
 
+  def to_proto(self, export_scope=None):
+    return self._primary_var.to_proto(export_scope=export_scope)
+
   @property
   def op(self):
     # We want cross-tower code that does some var.op.X calls
diff --git a/tensorflow/python/training/distribute.py b/tensorflow/python/training/distribute.py
@@ -1194,3 +1194,20 @@ def merge_fn(dist, vm):
 _default_tower_context = TowerContext(
     _default_distribution_strategy, tower_id=0)
 _default_tower_mode = _DefaultTowerThreadMode()
+
+
+# ------------------------------------------------------------------------------
+# We haven't yet implemented deserialization for DistributedVariables.
+# So here we catch any attempts to deserialize variables
+# when using distribution strategies.
+# pylint: disable=protected-access
+def _from_proto_fn(v, import_scope=None):
+  if has_distribution_strategy():
+    raise NotImplementedError(
+        "Deserialization of variables is not yet supported when using"
+        "distributed strategies.")
+  else:
+    resource_variable_ops._from_proto_fn(v, import_scope=import_scope)
+
+resource_variable_ops._from_proto_fn = _from_proto_fn
+# pylint: enable=protected-access
