Searched refs:run_in_function (Results 1 – 2 of 2) sorted by relevance
| /aosp_15_r20/external/tensorflow/tensorflow/python/distribute/ |
| H A D | strategy_test_lib.py | 64 def _maybe_run_in_function(fn, run_in_function=False): argument
65 if not run_in_function or not context.executing_eagerly():
449 def _test_numpy_dataset(self, strategy, session=None, run_in_function=False): argument
473 _maybe_run_in_function(lambda z: z, run_in_function), i)
629 def _test_run(self, strategy, run_in_function=False): argument
632 run_in_function))
636 lambda x: {"a": x * 2, "b": x * x}, run_in_function), args=(out1,))
642 lambda b, a: a + 2 * b + 2, run_in_function), kwargs=out2)
645 def _test_all_reduce_sum(self, strategy, run_in_function=False): argument
651 run_in_function=run_in_function)
[all …]
|
| H A D | tpu_strategy_test.py | 1320 self._test_numpy_dataset(strategy, run_in_function=True)
1328 self._test_run(strategy, run_in_function=True)
1336 self._test_all_reduce_sum(strategy, run_in_function=True)
1340 self._test_all_reduce_sum_gradients(strategy, run_in_function=True)
1344 self._test_all_reduce_sum_gradient_tape(strategy, run_in_function=True)
1348 self._test_all_reduce_mean(strategy, run_in_function=True)
1352 self._test_all_reduce_mean_gradients(strategy, run_in_function=True)
1356 self._test_all_reduce_mean_gradient_tape(strategy, run_in_function=True)
|