Home
last modified time | relevance | path

Searched refs:mesh_shape (Results 1 – 9 of 9) sorted by relevance

/external/tensorflow/tensorflow/python/tpu/
Dtopology.py52 def __init__(self, serialized=None, mesh_shape=None, device_coordinates=None): argument
82 self._mesh_shape = np.asarray(mesh_shape, dtype=np.int32)
108 self._mesh_shape = np.array(proto.mesh_shape, dtype=np.int32)
122 proto.mesh_shape))
128 proto.mesh_shape,
135 len(proto.mesh_shape)))
140 tasks = np.full(list(self.mesh_shape), -1, dtype=np.int32)
141 devices = np.full(list(self.mesh_shape), -1, dtype=np.int32)
150 def mesh_shape(self): member in Topology
229 proto.mesh_shape[:] = list(self._mesh_shape)
Dtopology_test.py31 mesh_shape=[1, 1, 1, 2],
40 original_topology.mesh_shape, new_topology.mesh_shape)
Ddevice_assignment.py392 topology_rank = len(topology.mesh_shape)
393 mesh_shape = topology.mesh_shape
422 if any(computation_footprint > mesh_shape):
425 computation_footprint, mesh_shape))
428 block_counts = mesh_shape // computation_footprint
437 mesh_shape))
/external/tensorflow/tensorflow/tools/api/golden/v2/
Dtensorflow.tpu.experimental.-topology.pbtxt14 name: "mesh_shape"
31 …argspec: "args=[\'self\', \'serialized\', \'mesh_shape\', \'device_coordinates\'], varargs=None, k…
/external/tensorflow/tensorflow/tools/api/golden/v1/
Dtensorflow.tpu.experimental.-topology.pbtxt14 name: "mesh_shape"
31 …argspec: "args=[\'self\', \'serialized\', \'mesh_shape\', \'device_coordinates\'], varargs=None, k…
/external/tensorflow/tensorflow/core/protobuf/tpu/
Dtopology.proto12 repeated int32 mesh_shape = 1; field
21 // [num_tasks, num_tpu_devices_per_task, len(mesh_shape)].
/external/tensorflow/tensorflow/compiler/mlir/tensorflow/utils/
Dtpu_rewrite_device_util.cc276 for (auto mesh_shape_dim : llvm::enumerate(topology_proto.mesh_shape())) in ParseTopologyAttr()
305 const int bound_x = topology_proto.mesh_shape(0); in ParseTopologyAttr()
306 const int bound_y = topology_proto.mesh_shape(1); in ParseTopologyAttr()
307 const int bound_z = topology_proto.mesh_shape(2); in ParseTopologyAttr()
308 const int bound_core = topology_proto.mesh_shape(3); in ParseTopologyAttr()
Dtpu_rewrite_device_util_test.cc114 std::string TopologyWithMeshShape(llvm::ArrayRef<int> mesh_shape) { in TopologyWithMeshShape() argument
116 for (int mesh_dim : mesh_shape) topology_proto.add_mesh_shape(mesh_dim); in TopologyWithMeshShape()
120 std::string TopologyWithMeshShapeAndTasks(llvm::ArrayRef<int> mesh_shape, in TopologyWithMeshShapeAndTasks() argument
124 for (int mesh_dim : mesh_shape) topology_proto.add_mesh_shape(mesh_dim); in TopologyWithMeshShapeAndTasks()
/external/tensorflow/tensorflow/compiler/mlir/tensorflow/tests/
Dtpu_rewrite.mlir261 // mesh_shape: 1
262 // mesh_shape: 1
263 // mesh_shape: 1
264 // mesh_shape: 2
1469 // mesh_shape: 1
1470 // mesh_shape: 2
1471 // mesh_shape: 1
1472 // mesh_shape: 2
1867 // mesh_shape: 2
1868 // mesh_shape: 1
[all …]