| /external/tensorflow/tensorflow/python/tpu/ |
| D | topology.py | 52 def __init__(self, serialized=None, mesh_shape=None, device_coordinates=None): argument
83 self._device_coordinates = np.asarray(device_coordinates, np.int32)
123 if len(proto.device_coordinates) != expected_coordinates_size:
129 len(proto.device_coordinates)))
131 coords = np.array(proto.device_coordinates, dtype=np.int32)
142 for task in xrange(self.device_coordinates.shape[0]):
143 for device in xrange(self.device_coordinates.shape[1]):
144 x, y, z, core = self.device_coordinates[task, device, :]
160 def device_coordinates(self): member in Topology
178 def task_ordinal_at_coordinates(self, device_coordinates): argument
[all …]
|
| D | topology_test.py | 32 device_coordinates=[[[0, 0, 0, 0], [0, 0, 0, 1]]],
42 original_topology.device_coordinates, new_topology.device_coordinates)
|
| D | device_assignment.py | 542 device_coordinates = topology.device_coordinates
553 replica_assignment.append(device_coordinates[task, device, :])
|
| /external/tensorflow/tensorflow/tools/api/golden/v2/ |
| D | tensorflow.tpu.experimental.-topology.pbtxt | 6 name: "device_coordinates"
31 …argspec: "args=[\'self\', \'serialized\', \'mesh_shape\', \'device_coordinates\'], varargs=None, k…
35 …argspec: "args=[\'self\', \'device_coordinates\', \'job\'], varargs=None, keywords=None, defaults=…
43 argspec: "args=[\'self\', \'device_coordinates\'], varargs=None, keywords=None, defaults=None"
47 …argspec: "args=[\'self\', \'device_coordinates\', \'job\'], varargs=None, keywords=None, defaults=…
51 argspec: "args=[\'self\', \'device_coordinates\'], varargs=None, keywords=None, defaults=None"
|
| /external/tensorflow/tensorflow/tools/api/golden/v1/ |
| D | tensorflow.tpu.experimental.-topology.pbtxt | 6 name: "device_coordinates"
31 …argspec: "args=[\'self\', \'serialized\', \'mesh_shape\', \'device_coordinates\'], varargs=None, k…
35 …argspec: "args=[\'self\', \'device_coordinates\', \'job\'], varargs=None, keywords=None, defaults=…
43 argspec: "args=[\'self\', \'device_coordinates\'], varargs=None, keywords=None, defaults=None"
47 …argspec: "args=[\'self\', \'device_coordinates\', \'job\'], varargs=None, keywords=None, defaults=…
51 argspec: "args=[\'self\', \'device_coordinates\'], varargs=None, keywords=None, defaults=None"
|
| /external/tensorflow/tensorflow/compiler/mlir/tensorflow/utils/ |
| D | tpu_rewrite_device_util.cc | 314 int x = topology_proto.device_coordinates(pos++); in ParseTopologyAttr()
315 int y = topology_proto.device_coordinates(pos++); in ParseTopologyAttr()
316 int z = topology_proto.device_coordinates(pos++); in ParseTopologyAttr()
317 int core = topology_proto.device_coordinates(pos++); in ParseTopologyAttr()
430 llvm::SmallVector<int64_t, 8> device_coordinates; in GetDeviceCoordinates() local
431 device_coordinates.reserve(device_assignment_attr.size()); in GetDeviceCoordinates()
443 device_coordinates.push_back(device_coordinate.getInt()); in GetDeviceCoordinates()
446 return device_coordinates; in GetDeviceCoordinates()
|
| D | tpu_rewrite_device_util_test.cc | 131 llvm::ArrayRef<int> device_coordinates) { in TopologyWithDeviceCoordinates() argument
139 for (int device_coordinate : device_coordinates) in TopologyWithDeviceCoordinates()
610 auto device_coordinates = status_or_device_coodinates.ConsumeValueOrDie(); in TEST() local
611 EXPECT_EQ(device_coordinates[0], 1); in TEST()
612 EXPECT_EQ(device_coordinates[1], 2); in TEST()
613 EXPECT_EQ(device_coordinates[2], 3); in TEST()
|
| /external/tensorflow/tensorflow/compiler/mlir/tensorflow/tests/ |
| D | tpu_rewrite.mlir | 267 // device_coordinates: 0
268 // device_coordinates: 0
269 // device_coordinates: 0
270 // device_coordinates: 0
271 // device_coordinates: 0
272 // device_coordinates: 0
273 // device_coordinates: 0
274 // device_coordinates: 1
1475 // device_coordinates: 0
1476 // device_coordinates: 0
[all …]
|
| /external/tensorflow/tensorflow/core/protobuf/tpu/ |
| D | topology.proto | 26 repeated int32 device_coordinates = 4; field
|
| /external/tensorflow/tensorflow/core/tpu/graph_rewrite/ |
| D | distributed_tpu_rewrite_pass.cc | 1417 int32 x = proto.device_coordinates(pos++); in ParseTopologyAttr()
1418 int32 y = proto.device_coordinates(pos++); in ParseTopologyAttr()
1419 int32 z = proto.device_coordinates(pos++); in ParseTopologyAttr()
1420 int32 core = proto.device_coordinates(pos++); in ParseTopologyAttr()
|