aboutsummaryrefslogtreecommitdiffhomepage
diff options
context:
space:
mode:
authorGravatar A. Unique TensorFlower <gardener@tensorflow.org>2018-03-26 20:02:24 -0700
committerGravatar TensorFlower Gardener <gardener@tensorflow.org>2018-03-26 20:06:29 -0700
commit8bcc574711b8770e8341f77d1a9b8370d72d7477 (patch)
tree074fa3435e40ac0826e68e8a42fa87e92dfabe5b
parent63cfd006fa1e848daeaf9ac74e2c9f8c42e401b1 (diff)
Include subscripts in the list of nodes accepted for replacement.
PiperOrigin-RevId: 190564824
-rw-r--r--tensorflow/contrib/autograph/pyct/templates.py12
1 files changed, 12 insertions, 0 deletions
diff --git a/tensorflow/contrib/autograph/pyct/templates.py b/tensorflow/contrib/autograph/pyct/templates.py
index fb99e0d4e5..baf7923fff 100644
--- a/tensorflow/contrib/autograph/pyct/templates.py
+++ b/tensorflow/contrib/autograph/pyct/templates.py
@@ -95,6 +95,15 @@ class ReplaceTransformer(gast.NodeTransformer):
self._check_inner_children_have_context(e)
for e in node.values:
self._check_inner_children_have_context(e)
+ elif isinstance(node, gast.Subscript):
+ self._check_inner_children_have_context(node.value)
+ self._check_inner_children_have_context(node.slice)
+ elif isinstance(node, gast.Slice):
+ self._check_inner_children_have_context(node.lower)
+ if node.upper:
+ self._check_inner_children_have_context(node.upper)
+ if node.step:
+ self._check_inner_children_have_context(node.step)
elif isinstance(node, gast.Name):
self._check_has_context(node)
elif isinstance(node, (gast.Str, gast.Num)):
@@ -127,6 +136,9 @@ class ReplaceTransformer(gast.NodeTransformer):
self._check_inner_children_have_context(e)
for e in node.values:
self._check_inner_children_have_context(e)
+ elif isinstance(node, gast.Subscript):
+ self._set_inner_child_context(node.value, ctx)
+ self._check_inner_children_have_context(node.slice)
elif isinstance(node, (gast.Str, gast.Num)):
pass
else: