MCPcopy Create free account
hub / github.com/DeepRec-AI/DeepRec / AddOp

Method AddOp

tensorflow/python/tpu/tpu.py:420–492  ·  view source on GitHub ↗
(self, op)

Source from the content-addressed store, hash-verified

418 return internal_control_inputs, external_control_inputs
419
420 def AddOp(self, op):
421 # pylint: disable=protected-access
422 if op.type in _BLACKLISTED_OPS:
423 logging.error("Operation of type %s (%s) is not supported on the TPU. "
424 "Execution will fail if this op is used in the graph. " %
425 (op.type, op.name))
426
427 if op.type in _UNSUPPORTED_OPS:
428 self._unsupported_ops.append(op)
429
430 if any(x.dtype._is_ref_dtype for x in op.inputs):
431 raise NotImplementedError(
432 "Non-resource Variables are not supported inside TPU computations "
433 "(operator name: %s)" % op.name)
434 if _TPU_REPLICATE_ATTR in op.node_def.attr:
435 raise ValueError("TPU computations cannot be nested")
436 op._set_attr(_TPU_REPLICATE_ATTR,
437 attr_value_pb2.AttrValue(s=self._name_as_bytes))
438 if self._outside_compilation_cluster:
439 op._set_attr(
440 _OUTSIDE_COMPILATION_ATTR,
441 attr_value_pb2.AttrValue(
442 s=compat.as_bytes(self._outside_compilation_cluster)))
443 if self._num_replicas > 1 or not self._outside_compilation_cluster:
444 # Prevent feeding or fetching anything that is being compiled,
445 # and any replicated outside_compilation Op.
446 op.graph.prevent_feeding(op)
447 op.graph.prevent_fetching(op)
448
449 # Remove any control edges from outer control flow contexts. These may cause
450 # mismatched frame errors.
451 (internal_control_inputs,
452 external_control_inputs) = self._RemoveExternalControlEdges(op)
453
454 if not op.inputs:
455 # Add a control edge from the control pivot to this op.
456 if not internal_control_inputs:
457 # pylint: disable=protected-access
458 op._add_control_input(self.GetControlPivot())
459 # pylint: enable=protected-access
460 else:
461 for index in xrange(len(op.inputs)):
462 x = op.inputs[index]
463 real_x = self.AddValue(x)
464 if real_x is not x:
465 op._update_input(index, real_x) # pylint: disable=protected-access
466
467 if external_control_inputs:
468 # Use an identity to pull control inputs as data inputs. Note that we
469 # ignore ops which don't have outputs. TODO(phawkins): fix that.
470 with ops.control_dependencies(None):
471 self.Enter()
472 external_control_inputs = [
473 array_ops.identity(x.outputs[0]).op
474 for x in external_control_inputs
475 if x.outputs
476 ]
477 self.Exit()

Callers 2

AddInnerOpMethod · 0.95

Calls 15

GetControlPivotMethod · 0.95
AddValueMethod · 0.95
EnterMethod · 0.95
anyFunction · 0.85
errorMethod · 0.80
_set_attrMethod · 0.80
prevent_feedingMethod · 0.80
prevent_fetchingMethod · 0.80
_add_control_inputMethod · 0.80
_update_inputMethod · 0.80
_add_control_inputsMethod · 0.80

Tested by

no test coverage detected