Searched refs:logical_device_id (Results 1 – 4 of 4) sorted by relevance
| /external/tensorflow/tensorflow/compiler/mlir/tensorflow/utils/ |
| D | xla_sharding_util.cc | 304 const int logical_device_id = sharding.tile_assignment_devices(0); in ExtractInputsForLogicalDevices() local
305 (*input_list)[logical_device_id].emplace_back(input_value); in ExtractInputsForLogicalDevices()
398 for (const auto logical_device_id : sharding.tile_assignment_devices()) { in GetTileShardedOutputsToMerge() local
400 sharding, logical_device_id, cluster_func_output_index); in GetTileShardedOutputsToMerge()
402 logical_device_id)[region_output_index]; in GetTileShardedOutputsToMerge()
599 int logical_device_id = 0; in RemapOutputsFromLogicalDevices() local
601 logical_device_id = output_sharding.tile_assignment_devices(0); in RemapOutputsFromLogicalDevices()
606 output_sharding_config, logical_device_id, output_index); in RemapOutputsFromLogicalDevices()
609 logical_device_id)[region_output_index]; in RemapOutputsFromLogicalDevices()
|
| /external/tensorflow/tensorflow/python/distribute/ |
| D | tpu_strategy.py | 404 def experimental_assign_to_logical_device(self, tensor, logical_device_id): argument
450 if (logical_device_id < 0 or
451 logical_device_id >= num_logical_devices_per_replica):
456 logical_device_id, num_logical_devices_per_replica))
458 tensor, logical_device_id, use_sharding_op=True)
1093 def experimental_logical_device(self, logical_device_id): argument
1096 if logical_device_id >= num_logical_devices_per_replica:
1100 logical_device_id, num_logical_devices_per_replica))
1102 self._logical_device_stack.append(logical_device_id)
1107 with ops.device(tpu.core(logical_device_id)):
[all …]
|
| /external/tensorflow/tensorflow/compiler/xla/python/ |
| D | xla.cc | 347 for (const auto& logical_device_id : span) { in PYBIND11_MODULE() local
349 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…
|