Searched refs:num_tpus_per_task (Results 1 – 7 of 7) sorted by relevance
| /external/tensorflow/tensorflow/compiler/mlir/tensorflow/utils/ |
| D | tpu_rewrite_device_util.cc | 191 const int num_tpus_per_task = tpu_devices[0].size(); in GetFullMeshTPUExecutionDeviceAssignment() local
192 const int num_tpu_devices = num_tasks * num_tpus_per_task; in GetFullMeshTPUExecutionDeviceAssignment()
206 const int task = i / num_tpus_per_task; in GetFullMeshTPUExecutionDeviceAssignment()
207 const int device = i % num_tpus_per_task; in GetFullMeshTPUExecutionDeviceAssignment()
265 llvm::StringRef topology_attr, int num_tasks, int num_tpus_per_task) { in ParseTopologyAttr() argument
288 if (topology_proto.num_tpu_devices_per_task() != num_tpus_per_task) in ParseTopologyAttr()
293 "), got ", num_tpus_per_task); in ParseTopologyAttr()
296 num_tasks * num_tpus_per_task * kTPUTopologyRank; in ParseTopologyAttr()
302 " (", num_tasks, " * ", num_tpus_per_task, " * ", kTPUTopologyRank, in ParseTopologyAttr()
313 for (int device = 0; device < num_tpus_per_task; ++device) { in ParseTopologyAttr()
[all …]
|
| /external/tensorflow/tensorflow/python/tpu/ |
| D | device_assignment.py | 540 computation_shape) <= topology.num_tasks * topology.num_tpus_per_task
550 task = logical_id // topology.num_tpus_per_task
551 device = logical_id % topology.num_tpus_per_task
|
| D | topology.py | 221 def num_tpus_per_task(self): member in Topology
|
| /external/tensorflow/tensorflow/tools/api/golden/v2/ |
| D | tensorflow.tpu.experimental.-topology.pbtxt | 26 name: "num_tpus_per_task"
|
| /external/tensorflow/tensorflow/tools/api/golden/v1/ |
| D | tensorflow.tpu.experimental.-topology.pbtxt | 26 name: "num_tpus_per_task"
|
| /external/tensorflow/tensorflow/core/tpu/graph_rewrite/ |
| D | distributed_tpu_rewrite_pass.cc | 1359 string* tpu_compilation_device, int* num_tpus_per_task, in GetTPUDeviceNames() argument
1377 replication_spec, device_set, num_tpus_per_task, tpu_devices)); in GetTPUDeviceNames()
1386 int num_tasks, int num_tpus_per_task, in ParseTopologyAttr() argument
1398 if (proto.num_tpu_devices_per_task() != num_tpus_per_task) { in ParseTopologyAttr()
1401 " != ", num_tpus_per_task, ")."); in ParseTopologyAttr()
1404 num_tasks * num_tpus_per_task * kTPUTopologyRank) { in ParseTopologyAttr()
1406 "device coordinates should be ", num_tasks, "x", num_tpus_per_task, "x", in ParseTopologyAttr()
1416 for (int device = 0; device < num_tpus_per_task; ++device) { in ParseTopologyAttr()
1504 int num_tasks, int num_tpus_per_task, in BuildFullMeshDeviceAssignment() argument
1508 int task = i / num_tpus_per_task; in BuildFullMeshDeviceAssignment()
[all …]
|
| D | distributed_tpu_rewrite_pass.h | 266 const tpu::TpuTopologyExternal& topology, int num_tpus_per_task,
|