1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
|
# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Wraps the body of a converted function with auxiliary constructs."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import gast
from tensorflow.python.autograph.core import converter
from tensorflow.python.autograph.pyct import templates
class FunctionBodyTransformer(converter.Base):
"""Wraps function bodies around autograph-specific boilerplate."""
def _name_for_current_scope(self):
innermost = self.enclosing_entities[-1]
if len(self.enclosing_entities) > 1:
parent = self.enclosing_entities[-2]
if isinstance(parent, gast.ClassDef):
# Methods also take the name of their class.
name = '%s/%s' % (parent.name, innermost.name)
else:
name = innermost.name
else:
name = innermost.name
# Sanitize the name.
# See https://www.tensorflow.org/api_docs/python/tf/Graph#name_scope
# TensorFlow doesn't like leading underscores at the top level.
while name[0] == '_':
name = name[1:]
return name
def visit_FunctionDef(self, node):
node = self.generic_visit(node)
final_body = []
indented_body = node.body
if node.body:
first_statement = node.body[0]
# Skip the docstring, if any.
if (isinstance(first_statement, gast.Expr) and
isinstance(first_statement.value, gast.Str)):
indented_body = indented_body[1:]
final_body.append(first_statement)
template = """
with ag__.function_scope(scope_name):
body
"""
scoped_body = templates.replace(
template,
scope_name=gast.Str(self._name_for_current_scope()),
body=indented_body)
final_body.extend(scoped_body)
node.body = final_body
return node
def transform(node, ctx):
return FunctionBodyTransformer(ctx).visit(node)
|
