Searched refs:device_dict (Results 1 – 4 of 4) sorted by relevance
60 device_dict = collections.defaultdict(list)108 device_dict[host_id].append(core_id)114 [len(core_ids) for core_ids in device_dict.values()])142 num_hosts=len(device_dict),
343 device_dict = {}349 if d_spec.job not in device_dict:350 device_dict[d_spec.job] = []353 while len(device_dict[d_spec.job]) <= d_spec.task:354 device_dict[d_spec.job].append([])356 device_dict[d_spec.job][d_spec.task].append(d)358 return device_dict385 device_dict = _group_device_list(devices)387 for _, devices_in_task in device_dict.items():487 device_dict = _group_device_list(devices)[all …]
138 def _verify_and_return_same_core_count(device_dict): argument141 {len(core_ids) for core_ids in device_dict.values()})144 'should never happen. Devices: {}'.format(device_dict))
606 device_dict, num_cores = TPUClusterResolver._get_device_dict_and_cores(609 self.assertEqual(device_dict, {})