From 6d893ecfb9ba2dfc3948215557d4f8ddaf7cf51b Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Thu, 6 Sep 2018 10:55:42 -0700
Subject: [PATCH] Ignore partitioned variable in TPU computation.

PiperOrigin-RevId: 211833891
---
 tensorflow/contrib/tpu/python/tpu/tpu.py | 15 +++++++++++++++
 1 file changed, 15 insertions(+)

diff --git a/tensorflow/contrib/tpu/python/tpu/tpu.py b/tensorflow/contrib/tpu/python/tpu/tpu.py
index 1e21cc52524..c1f90c3963d 100644
--- a/tensorflow/contrib/tpu/python/tpu/tpu.py
+++ b/tensorflow/contrib/tpu/python/tpu/tpu.py
@@ -652,13 +652,28 @@ def split_compile_and_replicate(computation,
       # TODO(phawkins): consider removing this code. It will
       # be less confusing to clients if they knowingly choose to use resource
       # variables.
+      # Partitioned variables is not supported (b/112311320).
+      def custom_getter(getter, name, *args, **kwargs):
+        partitioner = kwargs["partitioner"]
+        if partitioner is None:
+          return getter(name, *args, **kwargs)
+        else:
+          raise ValueError(
+              "Partitioned variables are not supported on TPU. Got "
+              "`partitioner` that is {}.".format(partitioner))
+
       vscope = variable_scope.get_variable_scope()
+
       saved_use_resource = vscope.use_resource
+      saved_custom_getter = vscope.custom_getter
+
       vscope.set_use_resource(True)
+      vscope.set_custom_getter(custom_getter)
 
       outputs = computation(*computation_inputs)
 
       vscope.set_use_resource(saved_use_resource)
+      vscope.set_custom_getter(saved_custom_getter)
 
     # If the computation returns `None`, make it an empty tuple.
     if outputs is None: