def node_to_graph(node, ctx, nocompile_decorators): """Convert Python code to equivalent TF graph mode code. Args: node: A Python AST node representing the code to convert. ctx: An EntityContext object. nocompile_decorators: A tuple containing decorators to be stripped from functions during conversion. Returns: A tuple (node, deps): * node: A Python ast node, representing the converted code. * deps: A set of strings, the fully qualified names of entity dependencies that this node has. """ # TODO(mdan): Verify arguments for correctness. # TODO(mdan): Factor out common elements. # These include: # * code move between blocks # * visiting blocks in transformers # Certain steps, especially canonicalization, insert new symbols into the # tree, which must be accounted. Although less efficient, it is most robust # to re-run the analysis. node = _static_analysis_pass(node, ctx) # Past this point, line numbers are no longer accurate so we ignore the # source. # TODO(mdan): Is it feasible to reconstruct intermediate source code? ctx.source_code = None node = decorators.transform(node, nocompile_decorators) node = break_statements.transform(node, ctx) node = asserts.transform(node, ctx) # Note: sequencing continue canonicalization before for loop one avoids # dealing with the extra loop increment operation that the for # canonicalization creates. node = continue_statements.transform(node, ctx) ctx.namespace['len'] = len node = _static_analysis_pass(node, ctx) node = for_loops.transform(node, ctx) # for_loops may insert new global references. node = builtin_functions.transform(node, ctx) # TODO(mdan): Kept for CL consistency. Remove. # builtin_functions may insert new global references. ctx.namespace['print'] = print node = _static_analysis_pass(node, ctx) node = call_trees.transform(node, ctx, config.DEFAULT_UNCOMPILED_MODULES, nocompile_decorators) node = control_flow.transform(node, ctx) # control_flow may create new symbols and change scopes. node = _static_analysis_pass(node, ctx) node = logical_expressions.transform(node) node = side_effect_guards.transform(node, ctx) return node
def node_to_graph(node, ctx, nocompile_decorators): """Convert Python code to equivalent TF graph mode code. Args: node: A Python AST node representing the code to convert. ctx: An EntityContext object. nocompile_decorators: A tuple containing decorators to be stripped from functions during conversion. Returns: A tuple (node, deps): * node: A Python ast node, representing the converted code. * deps: A set of strings, the fully qualified names of entity dependencies that this node has. """ # TODO(mdan): Verify arguments for correctness. # TODO(mdan): Factor out common elements. # These include: # * code move between blocks # * visiting blocks in transformers # Certain steps, especially canonicalization, insert new symbols into the # tree, which must be accounted. Although less efficient, it is most robust # to re-run the analysis. node = _static_analysis_pass(node, ctx) # Past this point, line numbers are no longer accurate so we ignore the # source. # TODO(mdan): Is it feasible to reconstruct intermediate source code? ctx.source_code = None node, deps = decorators.transform(node, nocompile_decorators) node = break_statements.transform(node, ctx) node = asserts.transform(node, ctx) # Note: sequencing continue canonicalization before for loop one avoids # dealing with the extra loop increment operation that the for # canonicalization creates. node = continue_statements.transform(node, ctx) ctx.namespace['len'] = len node = _static_analysis_pass(node, ctx) node = for_loops.transform(node, ctx) # for_loops may insert new global references. node = builtin_functions.transform(node, ctx) node = _static_analysis_pass(node, ctx) node = call_trees.transform(node, ctx, config.DEFAULT_UNCOMPILED_MODULES, nocompile_decorators) node = control_flow.transform(node, ctx) # control_flow may create new symbols and change scopes. node = _static_analysis_pass(node, ctx) node = logical_expressions.transform(node) node = side_effect_guards.transform(node, ctx) node = name_scopes.transform(node, ctx) return node, deps
def test_basic_for(self): def test_fn(l): s = 0 for e in l: s += e return s node = self.parse_and_analyze(test_fn, {}) node = for_loops.transform(node, self.ctx) with self.compiled(node) as result: l = [1, 2, 3] self.assertEqual(test_fn(l), result.test_fn(l)) l = [] self.assertEqual(test_fn(l), result.test_fn(l))
def test_basic_for(self): def test_fn(l): s = 0 for e in l: s += e return s node = self.parse_and_analyze(test_fn, {}) node = for_loops.transform(node, self.ctx) with self.compiled(node) as result: l = [1, 2, 3] self.assertEqual(test_fn(l), result.test_fn(l)) l = [] self.assertEqual(test_fn(l), result.test_fn(l))
def test_for_with_iterated_expression(self): eval_count = [0] def count_evals(x): eval_count[0] += 1 return x def test_fn(n): s = 0 for e in count_evals(range(n)): s += e return s node = self.parse_and_analyze(test_fn, {'count_evals': count_evals}) node = for_loops.transform(node, self.ctx) with self.compiled(node) as result: result.count_evals = count_evals self.assertEqual(test_fn(5), result.test_fn(5)) # count_evals ran twice, once for test_fn and another for result.test_fn self.assertEqual(eval_count[0], 2)
def test_for_with_iterated_expression(self): eval_count = [0] def count_evals(x): eval_count[0] += 1 return x def test_fn(n): s = 0 for e in count_evals(range(n)): s += e return s node = self.parse_and_analyze(test_fn, {'count_evals': count_evals}) node = for_loops.transform(node, self.ctx) with self.compiled(node) as result: result.count_evals = count_evals self.assertEqual(test_fn(5), result.test_fn(5)) # count_evals ran twice, once for test_fn and another for result.test_fn self.assertEqual(eval_count[0], 2)