Searched defs:num_tpus_per_task (Results 1 – 4 of 4) sorted by relevance
| /aosp_15_r20/external/tensorflow/tensorflow/compiler/mlir/tensorflow/utils/ |
| H A D | tpu_rewrite_device_util.cc | 187 const int num_tpus_per_task = tpu_devices[0].size(); in GetFullMeshTPUExecutionDeviceAssignment() local
261 llvm::StringRef topology_attr, int num_tasks, int num_tpus_per_task) { in ParseTopologyAttr()
348 const int num_tpus_per_task = tpu_devices[0].size(); in GetGeneralTPUExecutionDeviceAssignment() local
|
| /aosp_15_r20/external/tensorflow/tensorflow/python/tpu/ |
| H A D | topology.py | 225 def num_tpus_per_task(self): member in Topology
|
| /aosp_15_r20/external/tensorflow/tensorflow/core/tpu/graph_rewrite/ |
| H A D | update_tpu_embedding_ops_passes.cc | 98 int64 num_tpus_per_task = task_devices.size(); in Run() local
|
| H A D | distributed_tpu_rewrite_pass.cc | 1568 string* tpu_compilation_device, int* num_tpus_per_task, in GetTPUDeviceNames()
1595 int num_tasks, int num_tpus_per_task, in ParseTopologyAttr()
1714 int num_tasks, int num_tpus_per_task, in BuildFullMeshDeviceAssignment()
1776 const tpu::TpuTopologyExternal& tpu_topology, int num_tpus_per_task, in BuildDeviceAssignment()
4388 int num_tpus_per_task; in GetDeviceTopology() local
|