Processes the newly-added TF_Operations in `graph`.
(graph)
| 235 | |
| 236 | |
| 237 | def _ProcessNewOps(graph): |
| 238 | """Processes the newly-added TF_Operations in `graph`.""" |
| 239 | # Maps from a node to the names of the ops it's colocated with, if colocation |
| 240 | # is specified in the attributes. |
| 241 | colocation_pairs = {} |
| 242 | |
| 243 | for new_op in graph._add_new_tf_operations(compute_devices=False): # pylint: disable=protected-access |
| 244 | original_device = new_op.device |
| 245 | new_op._set_device('') # pylint: disable=protected-access |
| 246 | colocation_names = _GetColocationNames(new_op) |
| 247 | if colocation_names: |
| 248 | colocation_pairs[new_op] = colocation_names |
| 249 | # Don't set a device for this op, since colocation constraints override |
| 250 | # device functions and the original device. Note that this op's device may |
| 251 | # still be set by the loop below. |
| 252 | # TODO(skyewm): why does it override the original device? |
| 253 | else: |
| 254 | with _MaybeDevice(original_device): |
| 255 | graph._apply_device_functions(new_op) # pylint: disable=protected-access |
| 256 | |
| 257 | # The following loop populates the device field of ops that are colocated |
| 258 | # with another op. This is implied by the colocation attribute, but we |
| 259 | # propagate the device field for completeness. |
| 260 | for op, coloc_op_list in colocation_pairs.items(): |
| 261 | coloc_device = None |
| 262 | # Find any device in the list of colocated ops that have a device, if it |
| 263 | # exists. We assume that if multiple ops have devices, they refer to the |
| 264 | # same device. Otherwise, a runtime error will occur since the colocation |
| 265 | # property cannot be guaranteed. Note in TF2 colocations have been removed |
| 266 | # from the public API and will be considered a hint, so there is no runtime |
| 267 | # error. |
| 268 | # |
| 269 | # One possible improvement is to try to check for compatibility of all |
| 270 | # devices in this list at import time here, which would require |
| 271 | # implementing a compatibility function for device specs in python. |
| 272 | for coloc_op_name in coloc_op_list: |
| 273 | try: |
| 274 | coloc_op = graph._get_operation_by_name_unsafe(coloc_op_name) # pylint: disable=protected-access |
| 275 | except KeyError: |
| 276 | # Do not error in TF2 if the colocation cannot be guaranteed |
| 277 | if tf2.enabled() or control_flow_util.EnableControlFlowV2(graph): |
| 278 | continue |
| 279 | |
| 280 | raise ValueError('Specified colocation to an op that ' |
| 281 | 'does not exist during import: %s in %s' % |
| 282 | (coloc_op_name, op.name)) |
| 283 | if coloc_op.device: |
| 284 | coloc_device = pydev.DeviceSpec.from_string(coloc_op.device) |
| 285 | break |
| 286 | if coloc_device: |
| 287 | op._set_device(coloc_device) # pylint: disable=protected-access |
| 288 | |
| 289 | |
| 290 | def _GetColocationNames(op): |
no test coverage detected