Searched refs:kTPUTopologyRank (Results 1 – 2 of 2) sorted by relevance
| /external/tensorflow/tensorflow/compiler/mlir/tensorflow/utils/ |
| D | tpu_rewrite_device_util.cc | 51 constexpr int kTPUTopologyRank = 4; variable
271 if (topology_proto.mesh_shape_size() != kTPUTopologyRank) in ParseTopologyAttr()
273 "'", kTopologyAttr, "' 'mesh_shape' must be rank ", kTPUTopologyRank, in ParseTopologyAttr()
296 num_tasks * num_tpus_per_task * kTPUTopologyRank; in ParseTopologyAttr()
301 "' must be 'num_tasks' * 'num_tpus_per_task' * ", kTPUTopologyRank, in ParseTopologyAttr()
302 " (", num_tasks, " * ", num_tpus_per_task, " * ", kTPUTopologyRank, in ParseTopologyAttr()
358 num_replicas * num_cores_per_replica * kTPUTopologyRank; in GetGeneralTPUExecutionDeviceAssignment()
364 kTPUTopologyRank, " (", num_replicas, " * ", num_cores_per_replica, in GetGeneralTPUExecutionDeviceAssignment()
365 " * ", kTPUTopologyRank, "), got ", device_assignment_attr.size()); in GetGeneralTPUExecutionDeviceAssignment()
|
| /external/tensorflow/tensorflow/core/tpu/graph_rewrite/ |
| D | distributed_tpu_rewrite_pass.cc | 82 constexpr int kTPUTopologyRank = 4; variable
1388 static_assert(4 == kTPUTopologyRank, "Assumes the topology rank is 4"); in ParseTopologyAttr()
1391 if (proto.mesh_shape_size() != kTPUTopologyRank) { in ParseTopologyAttr()
1393 kTPUTopologyRank); in ParseTopologyAttr()
1404 num_tasks * num_tpus_per_task * kTPUTopologyRank) { in ParseTopologyAttr()
1407 kTPUTopologyRank, "; got ", proto.device_coordinates_size()); in ParseTopologyAttr()
1446 static_assert(4 == kTPUTopologyRank, "Assumes the topology rank is 4"); in ParseDeviceAssignmentAttr()
1449 num_replicas * num_cores_per_replica * kTPUTopologyRank; in ParseDeviceAssignmentAttr()
1454 ") * ", kTPUTopologyRank, " got ", device_assignment_attr.size()); in ParseDeviceAssignmentAttr()
|