Searched defs:logical_device_id (Results 1 – 4 of 4) sorted by relevance
298 const int logical_device_id = sharding.tile_assignment_devices(0); in ExtractInputsForLogicalDevices() local400 for (const auto logical_device_id : sharding.tile_assignment_devices()) { in GetTileShardedOutputsToMerge() local625 int logical_device_id = 0; in RemapOutputsFromLogicalDevices() local
441 def experimental_assign_to_logical_device(self, tensor, logical_device_id): argument1151 def experimental_logical_device(self, logical_device_id): argument1711 def experimental_logical_device(self, logical_device_id): argument
368 for (const auto& logical_device_id : span) { in PYBIND11_MODULE() local
842 int32_t logical_device_id = in __xla_cpu_runtime_CollectivePermute() local