Home
last modified time | relevance | path

Searched refs:tpu_function (Results 1 – 9 of 9) sorted by relevance

/external/tensorflow/tensorflow/python/tpu/
Dtraining_loop.py30 from tensorflow.python.tpu import tpu_function
119 number_of_shards = tpu_function.get_tpu_context().number_of_shards
172 num_replicas = tpu_function.get_tpu_context().number_of_shards
Dtpu_optimizer.py26 from tensorflow.python.tpu import tpu_function
144 num_shards = tpu_function.get_tpu_context().number_of_shards
Dfeature_column_v2_test.py39 from tensorflow.python.tpu import tpu_function
249 with tpu_function.tpu_shard_context(1):
276 with tpu_function.tpu_shard_context(1):
607 with tpu_function.tpu_shard_context(1):
694 with tpu_function.tpu_shard_context(1):
DBUILD50 ":tpu_function",
232 ":tpu_function",
272 ":tpu_function",
320 name = "tpu_function",
321 srcs = ["tpu_function.py"],
Dfeature_column.py29 from tensorflow.python.tpu import tpu_function
675 return tpu_function.get_tpu_context().number_of_shards is None
Dtpu.py57 from tensorflow.python.tpu import tpu_function
1403 with tpu_function.tpu_shard_context(
/external/tensorflow/tensorflow/python/tpu/ops/
Dtpu_ops.py30 from tensorflow.python.tpu import tpu_function
35 num_shards = tpu_function.get_tpu_context().number_of_shards
/external/tensorflow/tensorflow/python/distribute/
Dtpu_strategy.py1438 def tpu_function(args, kwargs): function
1510 tpu_function = def_function.function(tpu_function)
1511 self._tpu_function_cache[fn] = tpu_function
1512 return tpu_function
Dtpu_strategy_test.py707 def tpu_function(sparse): function
716 strategy.run(tpu_function, args=(next(iterator),)))
757 def tpu_function(sparse): function
766 strategy.run(tpu_function, args=(next(iterator),)))
810 def tpu_function(sparse): function
816 strategy.run(tpu_function, args=(next(iterator),)))