Searched refs:logical_device_id (Results 1 – 6 of 6) sorted by relevance
| /external/tensorflow/tensorflow/compiler/mlir/tensorflow/utils/ |
| D | xla_sharding_util.cc | 298 const int logical_device_id = sharding.tile_assignment_devices(0); in ExtractInputsForLogicalDevices() local
299 (*input_list)[logical_device_id].emplace_back(input_value); in ExtractInputsForLogicalDevices()
400 for (const auto logical_device_id : sharding.tile_assignment_devices()) { in GetTileShardedOutputsToMerge() local
403 location, cluster_to_core_index, logical_device_id, in GetTileShardedOutputsToMerge()
408 cluster_idx + logical_device_id)[region_output_index]; in GetTileShardedOutputsToMerge()
625 int logical_device_id = 0; in RemapOutputsFromLogicalDevices() local
627 logical_device_id = output_sharding.tile_assignment_devices(0); in RemapOutputsFromLogicalDevices()
633 logical_device_id, output_index, in RemapOutputsFromLogicalDevices()
639 cluster_idx + logical_device_id)[region_output_index]; in RemapOutputsFromLogicalDevices()
|
| /external/tensorflow/tensorflow/python/distribute/ |
| D | tpu_strategy.py | 441 def experimental_assign_to_logical_device(self, tensor, logical_device_id): argument
494 if (logical_device_id < 0 or
495 logical_device_id >= num_logical_devices_per_replica):
500 logical_device_id, num_logical_devices_per_replica))
502 tensor, logical_device_id, use_sharding_op=True)
1151 def experimental_logical_device(self, logical_device_id): argument
1154 if logical_device_id >= num_logical_devices_per_replica:
1158 logical_device_id, num_logical_devices_per_replica))
1160 self._logical_device_stack.append(logical_device_id)
1165 with ops.device(tpu.core(logical_device_id)):
[all …]
|
| /external/tensorflow/tensorflow/compiler/xla/python/ |
| D | xla.cc | 368 for (const auto& logical_device_id : span) { in PYBIND11_MODULE() local
370 logical_device_id.replica, logical_device_id.partition)); in PYBIND11_MODULE()
|
| /external/tensorflow/tensorflow/tools/api/golden/v2/ |
| D | tensorflow.distribute.-t-p-u-strategy.pbtxt | 37 …argspec: "args=[\'self\', \'tensor\', \'logical_device_id\'], varargs=None, keywords=None, default…
|
| /external/pytorch/torch/fx/experimental/ |
| D | accelerator_partitioner.py | 341 total_size_of_graph, logical_device_id=device_with_max_mem.logical_id
386 self, total_size_of_graph, logical_device_id: int = 0
397 partition_0.logical_device_ids = [logical_device_id]
|
| /external/tensorflow/tensorflow/compiler/xla/service/cpu/ |
| D | cpu_runtime.cc | 842 int32_t logical_device_id = in __xla_cpu_runtime_CollectivePermute() local
851 if (from == logical_device_id) { in __xla_cpu_runtime_CollectivePermute()
861 participant.replica_id = logical_device_id; in __xla_cpu_runtime_CollectivePermute()
|