Function _preprocess_grad

tensorflow/python/ops/while_v2.py:475–505

Returns the initial gradient to be used for a given output tensor. Args: grad: the original gradient Tensor passed to the gradient function. body_graph_output: the corresponding Tensor in the body graph. while_op_output: the corresponding Tensor output of the While op. Returns:

(grad, body_graph_output, while_op_output)

def _preprocess_grad(grad, body_graph_output, while_op_output):
  """Returns the initial gradient to be used for a given output tensor.

  Args:
    grad: the original gradient Tensor passed to the gradient function.
    body_graph_output: the corresponding Tensor in the body graph.
    while_op_output: the corresponding Tensor output of the While op.

  Returns:
    A Tensor or None.
  """
  # Set the incoming gradient of non-trainable inputs to None. It is possible
  # that we receive non-None gradients for non-trainable types in nested while
  # loops because we accumulate outputs of the inner while as variant tensors
  # which are trainable and hence receive zeros_like tensors in the gradient
  # pass. The non-trainable tensors then receive the popped zeros tensor from
  # this zeros variant. The gradient for the loop vars corresponding to these
  # tensors is None or zeros (this happens only if the loop var is accumulated
  # as well) in _grad_fn so we reset these.
  # TODO(b/118712257): Remove once we can handle None output grads in _grad_fn.
  if not _is_trainable(body_graph_output):
    return None

  # GradientTape initializes resource and variant grads as None instead of
  # zeros. Set to zeros so _GradientsHelper computes the gradients instead of
  # returning None.
  if (while_op_output.dtype in (dtypes.resource, dtypes.variant)
      and grad is None):
    return _zeros_like(while_op_output)

  return grad


# TODO(skyewm): make this return constants if op_output's shape is fully