Home
last modified time | relevance | path

Searched refs:topology_proto (Results 1 – 3 of 3) sorted by relevance

/external/tensorflow/tensorflow/compiler/mlir/tensorflow/utils/
Dtpu_rewrite_device_util_test.cc115 tpu::TopologyProto topology_proto; in TopologyWithMeshShape() local
116 for (int mesh_dim : mesh_shape) topology_proto.add_mesh_shape(mesh_dim); in TopologyWithMeshShape()
117 return topology_proto.SerializeAsString(); in TopologyWithMeshShape()
123 tpu::TopologyProto topology_proto; in TopologyWithMeshShapeAndTasks() local
124 for (int mesh_dim : mesh_shape) topology_proto.add_mesh_shape(mesh_dim); in TopologyWithMeshShapeAndTasks()
125 topology_proto.set_num_tasks(num_tasks); in TopologyWithMeshShapeAndTasks()
126 topology_proto.set_num_tpu_devices_per_task(num_tpu_devices_per_task); in TopologyWithMeshShapeAndTasks()
127 return topology_proto.SerializeAsString(); in TopologyWithMeshShapeAndTasks()
132 tpu::TopologyProto topology_proto; in TopologyWithDeviceCoordinates() local
133 topology_proto.add_mesh_shape(2); in TopologyWithDeviceCoordinates()
[all …]
Dtpu_rewrite_device_util.cc266 tpu::TopologyProto topology_proto; in ParseTopologyAttr() local
267 if (!topology_proto.ParseFromString(topology_attr.str())) in ParseTopologyAttr()
271 if (topology_proto.mesh_shape_size() != kTPUTopologyRank) in ParseTopologyAttr()
274 ", got rank ", topology_proto.mesh_shape_size()); in ParseTopologyAttr()
276 for (auto mesh_shape_dim : llvm::enumerate(topology_proto.mesh_shape())) in ParseTopologyAttr()
283 if (topology_proto.num_tasks() != num_tasks) in ParseTopologyAttr()
286 kTopologyAttr, "' (", topology_proto.num_tasks(), "), got ", num_tasks); in ParseTopologyAttr()
288 if (topology_proto.num_tpu_devices_per_task() != num_tpus_per_task) in ParseTopologyAttr()
292 kTopologyAttr, "' (", topology_proto.num_tpu_devices_per_task(), in ParseTopologyAttr()
297 if (topology_proto.device_coordinates_size() != in ParseTopologyAttr()
[all …]
/external/tensorflow/tensorflow/core/protobuf/tpu/
DBUILD45 name = "topology_proto",
91 proto_lib = ":topology_proto",
122 # deps = [":topology_proto"],