def _gradient_function(op_name, attr_tuple, num_inputs, inputs, outputs,
out_grads, skip_input_indices, forward_pass_name_scope):
"""Calls the gradient function of the op.
Args:
op_name: the name of the op to be differentiated.
attr_tuple: the attrs, as a tuple.
num_inputs: the number of inputs to the op.
inputs: inputs to the original operation.
outputs: outputs to the original operation.
out_grads: gradients of the operation wrt its outputs.
skip_input_indices: a tuple that is passed to the gradient function,
indicating which inputs to skip calculating the gradient for
forward_pass_name_scope: the namescope of the op in the forward pass.
Returns:
The gradients with respect to the inputs of the function, as a list.
"""
mock_op = _MockOp(attr_tuple, inputs, outputs, op_name, skip_input_indices)
grad_fn = ops._gradient_registry.lookup(op_name) # pylint: disable=protected-access
if grad_fn is None:
return [None] * num_inputs
# This does not work with v1 TensorArrays.
if ops.executing_eagerly_outside_functions(
) or control_flow_util.EnableControlFlowV2(ops.get_default_graph()):
gradient_name_scope = "gradient_tape/"
if forward_pass_name_scope:
gradient_name_scope += forward_pass_name_scope + "/"
with ops.name_scope(gradient_name_scope):
return grad_fn(mock_op, *out_grads)
else:
return grad_fn(mock_op, *out_grads)
def _assign_moving_average(self, variable, value, momentum):
with ops.name_scope(None, 'AssignMovingAvg',
[variable, value, momentum]) as scope:
# TODO(b/120571621): We want to avoid colocating the variables here
# since TPUStrategy does not implement replica local variables.
# Remove this hack once we support TPULocalVariables.
is_tpu_strategy = False
if distribution_strategy_context.has_distribution_strategy():
distribute = distribution_strategy_context.get_distribution_strategy(
)
if distribute.__class__.__name__ == 'TPUStrategy':
is_tpu_strategy = True
# TODO(apassos,srbs,skyewm): the colocation constraints here are disabled
# because of a bug which leads cond_v2/while_v2 to skip rewriting them
# creating conflicts.
if (control_flow_util.EnableControlFlowV2(ops.get_default_graph())
or is_tpu_strategy):
cm = contextlib.contextmanager(lambda: (yield))()
else:
cm = ops.colocate_with(variable)
with cm:
decay = ops.convert_to_tensor(1.0 - momentum, name='decay')
if decay.dtype != variable.dtype.base_dtype:
decay = math_ops.cast(decay, variable.dtype.base_dtype)
update_delta = (variable -
math_ops.cast(value, variable.dtype)) * decay
return state_ops.assign_sub(variable, update_delta, name=scope)
def _ProcessNewOps(graph):
"""Processes the newly-added TF_Operations in `graph`."""
# Maps from a node to the names of the ops it's colocated with, if colocation
# is specified in the attributes.
colocation_pairs = {}
for new_op in graph._add_new_tf_operations(compute_devices=False): # pylint: disable=protected-access
original_device = new_op.device
new_op._set_device('') # pylint: disable=protected-access
colocation_names = _GetColocationNames(new_op)
if colocation_names:
colocation_pairs[new_op] = colocation_names
# Don't set a device for this op, since colocation constraints override
# device functions and the original device. Note that this op's device may
# still be set by the loop below.
# TODO(skyewm): why does it override the original device?
else:
with _MaybeDevice(original_device):
graph._apply_device_functions(new_op) # pylint: disable=protected-access
# The following loop populates the device field of ops that are colocated
# with another op. This is implied by the colocation attribute, but we
# propagate the device field for completeness.
for op, coloc_op_list in colocation_pairs.items():
coloc_device = None
# Find any device in the list of colocated ops that have a device, if it
# exists. We assume that if multiple ops have devices, they refer to the
# same device. Otherwise, a runtime error will occur since the colocation
# property cannot be guaranteed. Note in TF2 colocations have been removed
# from the public API and will be considered a hint, so there is no runtime
# error.
#
# One possible improvement is to try to check for compatibility of all
# devices in this list at import time here, which would require
# implementing a compatibility function for device specs in python.
for coloc_op_name in coloc_op_list:
try:
coloc_op = graph._get_operation_by_name_unsafe(coloc_op_name) # pylint: disable=protected-access
except KeyError:
# Do not error in TF2 if the colocation cannot be guaranteed
if tf2.enabled() or control_flow_util.EnableControlFlowV2(
graph):
continue
raise ValueError('Specified colocation to an op that '
'does not exist during import: %s in %s' %
(coloc_op_name, op.name))
if coloc_op.device:
coloc_device = pydev.DeviceSpec.from_string(coloc_op.device)
break
if coloc_device:
op._set_device(coloc_device) # pylint: disable=protected-access
def control_flow_v2_enabled(): # pylint: disable=invalid-name
"""Returns `True` if v2 control flow is enabled.
Note: v2 control flow is always enabled inside of tf.function.
"""
return control_flow_util.EnableControlFlowV2(ops.get_default_graph())
def __init__(self,
dtype,
size=None,
dynamic_size=None,
clear_after_read=None,
tensor_array_name=None,
handle=None,
flow=None,
infer_shape=True,
element_shape=None,
colocate_with_first_write_call=True,
name=None):
"""Construct a new TensorArray or wrap an existing TensorArray handle.
A note about the parameter `name`:
The name of the `TensorArray` (even if passed in) is uniquified: each time
a new `TensorArray` is created at runtime it is assigned its own name for
the duration of the run. This avoids name collisions if a `TensorArray`
is created within a `while_loop`.
Args:
dtype: (required) data type of the TensorArray.
size: (optional) int32 scalar `Tensor`: the size of the TensorArray.
Required if handle is not provided.
dynamic_size: (optional) Python bool: If true, writes to the TensorArray
can grow the TensorArray past its initial size. Default: False.
clear_after_read: Boolean (optional, default: True). If True, clear
TensorArray values after reading them. This disables read-many
semantics, but allows early release of memory.
tensor_array_name: (optional) Python string: the name of the TensorArray.
This is used when creating the TensorArray handle. If this value is
set, handle should be None.
handle: (optional) A `Tensor` handle to an existing TensorArray. If this
is set, tensor_array_name should be None. Only supported in graph mode.
flow: (optional) A float `Tensor` scalar coming from an existing
`TensorArray.flow`. Only supported in graph mode.
infer_shape: (optional, default: True) If True, shape inference
is enabled. In this case, all elements must have the same shape.
element_shape: (optional, default: None) A `TensorShape` object specifying
the shape constraints of each of the elements of the TensorArray.
Need not be fully defined.
colocate_with_first_write_call: If `True`, the TensorArray will be
colocated on the same device as the Tensor used on its first write
(write operations include `write`, `unstack`, and `split`). If `False`,
the TensorArray will be placed on the device determined by the
device context available during its initialization.
name: A name for the operation (optional).
Raises:
ValueError: if both handle and tensor_array_name are provided.
TypeError: if handle is provided but is not a Tensor.
"""
if context.executing_eagerly():
implementation = _EagerTensorArray
else:
if control_flow_util.EnableControlFlowV2(ops.get_default_graph()):
implementation = _GraphTensorArrayV2
else:
implementation = _GraphTensorArray
self._implementation = implementation(
dtype,
size=size,
dynamic_size=dynamic_size,
clear_after_read=clear_after_read,
tensor_array_name=tensor_array_name,
handle=handle,
flow=flow,
infer_shape=infer_shape,
element_shape=element_shape,
colocate_with_first_write_call=colocate_with_first_write_call,
name=name)
self._implementation.parent = weakref.ref(self)
