Home
last modified time | relevance | path

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

/aosp_15_r20/external/tensorflow/tensorflow/dtensor/mlir/
H A Dmove_compilation_to_host.cc223 mlir::func::FuncOp tpu_function = xla_mesh->second.func(); in HandleCompilationOps() local
254 tpu_function.walk( in HandleCompilationOps()
259 return tpu_function.emitOpError(llvm::formatv( in HandleCompilationOps()
329 builder.setInsertionPointToStart(&tpu_function.front()); in HandleCompilationOps()
332 xla_mesh->first, tpu_function, device_insertion_point, nullptr}); in HandleCompilationOps()
344 tpu_function.walk([&](mlir::Operation* op) { in HandleCompilationOps()
/aosp_15_r20/external/tensorflow/tensorflow/python/tpu/
H A Dtraining_loop.py26 from tensorflow.python.tpu import tpu_function
111 number_of_shards = tpu_function.get_tpu_context().number_of_shards
164 num_replicas = tpu_function.get_tpu_context().number_of_shards
H A Dfeature_column_v2_test.py37 from tensorflow.python.tpu import tpu_function
247 with tpu_function.tpu_shard_context(1):
274 with tpu_function.tpu_shard_context(1):
625 with tpu_function.tpu_shard_context(1):
712 with tpu_function.tpu_shard_context(1):
H A Dtpu_optimizer.py22 from tensorflow.python.tpu import tpu_function
145 num_shards = tpu_function.get_tpu_context().number_of_shards
H A DBUILD232 ":tpu_function",
271 ":tpu_function",
317 name = "tpu_function",
318 srcs = ["tpu_function.py"],
H A Dfeature_column.py25 from tensorflow.python.tpu import tpu_function
673 return tpu_function.get_tpu_context().number_of_shards is None
H A Dtpu.py53 from tensorflow.python.tpu import tpu_function
1440 with tpu_function.tpu_shard_context(
/aosp_15_r20/external/tensorflow/tensorflow/python/distribute/
H A Dtpu_strategy_test.py874 def tpu_function(sparse): function
883 strategy.run(tpu_function, args=(next(iterator),)))
924 def tpu_function(sparse): function
933 strategy.run(tpu_function, args=(next(iterator),)))
977 def tpu_function(sparse): function
983 strategy.run(tpu_function, args=(next(iterator),)))
1053 def tpu_function(composite): function
1060 strategy.run(tpu_function, args=(test_composite,)))
H A Dtpu_strategy.py1569 def tpu_function(args, kwargs): function
1648 tpu_function = def_function.function(tpu_function)
1649 self._tpu_function_cache[fn] = tpu_function
1650 return tpu_function
/aosp_15_r20/external/tensorflow/tensorflow/python/tpu/ops/
H A Dtpu_ops.py25 from tensorflow.python.tpu import tpu_function
30 num_shards = tpu_function.get_tpu_context().number_of_shards
H A DBUILD17 "//tensorflow/python/tpu:tpu_function",