Fix call count test
Due to relying on mutations to a Python nonlocal
variable to make assertions about the call count,
these tests were actually tracking the number of
times the Python function was called to create the
TFF AST, not the number of times the computation
was run.
@@ -582,41 +582,20 @@ class FederatedComputationsTest(parameterized.TestCase):
class ComputationsTest(parameterized.TestCase):
@core_test.executors
- def test_deterministic_computation_called_twice_is_invoked_once(self):
- count = 0
-
- @tff.tf_computation
- def foo():
- nonlocal count
- count += 1
- return 10
-
- @tff.federated_computation
- def bar():
- return foo(), foo()
-
- bar()
- self.assertEqual(count, 1)
-
- @core_test.executors
- def test_nondeterministic_computation_called_twice_is_invoked_twice(self):
+ def test_tf_computation_called_twice_is_invoked_twice(self):
self.skipTest(
'b/139135080: Recognize distinct instantiations of the same TF code as '
'(potentially) distinct at construction time.')
- count = 0
-
@tff.tf_computation
- def foo():
- nonlocal count
- count += 1
+ def get_random():
return tf.random.normal([])
@tff.federated_computation
- def bar():
- return foo(), foo()
+ def get_two_random():
+ return get_random(), get_random()
- bar()
- self.assertEqual(count, 2)
+ first_random, second_random = get_two_random()
+ self.assertNotEqual(first_random, second_random)
if __name__ == '__main__':
