From d4c014f6fecc791eba0436fb500122d354354252 Mon Sep 17 00:00:00 2001 From: "A. Unique TensorFlower" Date: Fri, 13 Dec 2019 17:41:15 -0800 Subject: [PATCH] Add test for creating variable inside function passed to tf.vectorized_map. Switch pfor logic to use tf.function. PiperOrigin-RevId: 285504833 Change-Id: I1d0712349dae0cbe49a05684a0bdbba48eb517ff --- .../ops/parallel_for/control_flow_ops.py | 4 +-- .../ops/parallel_for/control_flow_ops_test.py | 31 ------------------- 2 files changed, 2 insertions(+), 33 deletions(-) diff --git a/tensorflow/python/ops/parallel_for/control_flow_ops.py b/tensorflow/python/ops/parallel_for/control_flow_ops.py index ff6701142da..a1c7588ea8f 100644 --- a/tensorflow/python/ops/parallel_for/control_flow_ops.py +++ b/tensorflow/python/ops/parallel_for/control_flow_ops.py @@ -22,7 +22,7 @@ from __future__ import print_function import functools from tensorflow.python.eager import context -from tensorflow.python.eager import def_function +from tensorflow.python.eager import function from tensorflow.python.framework import indexed_slices from tensorflow.python.framework import ops from tensorflow.python.framework import sparse_tensor @@ -185,7 +185,7 @@ def pfor(loop_fn, iters, parallel_iterations=None): # XLA compilation. The latter is so that we don't compile operations like # tf.placeholder that are created by the loop body. if context.executing_eagerly() or _is_under_xla_context(): - f = def_function.function(f) + f = function.defun(f) return f() diff --git a/tensorflow/python/ops/parallel_for/control_flow_ops_test.py b/tensorflow/python/ops/parallel_for/control_flow_ops_test.py index 828a9e018c1..d016393a39e 100644 --- a/tensorflow/python/ops/parallel_for/control_flow_ops_test.py +++ b/tensorflow/python/ops/parallel_for/control_flow_ops_test.py @@ -1477,36 +1477,5 @@ class PartitionedCallTest(PForTestCase): self._test_loop_fn(loop_fn, 4) -class VariableTest(PForTestCase): - - def test_create_variable_once(self): - x = array_ops.ones(shape=(3, 2, 2), dtype=dtypes.float32) - y = array_ops.ones(shape=(2, 3), dtype=dtypes.float32) - a_var = [] - - def f(z): - if not a_var: - a_var.append(variables.Variable(lambda: y, name="a")) - return math_ops.matmul(z, a_var[0] / 16) - - pfor_control_flow_ops.vectorized_map(f, x) - - @test_util.run_v2_only - def test_create_variable_repeated(self): - x = array_ops.ones(shape=(3, 2, 2), dtype=dtypes.float32) - y = array_ops.ones(shape=(2, 3), dtype=dtypes.float32) - - def f(z): - a_var = variables.Variable(lambda: y, name="a") / 4 - return math_ops.matmul(z, a_var / 16) - - # Note that this error is only raised under v2 behavior. - with self.assertRaisesRegexp( - ValueError, - "tf.function-decorated function tried to create variables on non-first" - ): - pfor_control_flow_ops.vectorized_map(f, x) - - if __name__ == "__main__": test.main()