Searched defs:logical_device_id (Results 1 – 3 of 3) sorted by relevance
302 const int logical_device_id = sharding.tile_assignment_devices(0); in ExtractInputsForLogicalDevices() local399 for (const auto logical_device_id : sharding.tile_assignment_devices()) { in GetTileShardedOutputsToMerge() local598 int logical_device_id = 0; in RemapOutputsFromLogicalDevices() local
405 def experimental_assign_to_logical_device(self, tensor, logical_device_id): argument1094 def experimental_logical_device(self, logical_device_id): argument1561 def experimental_logical_device(self, logical_device_id): argument
316 for (const auto& logical_device_id : span) { in PYBIND11_MODULE() local