Function DynamicStitchGrad

tensorflow/cc/gradients/data_flow_grad.cc:109–149

REGISTER_GRADIENT_OP("DynamicPartition", DynamicPartitionGrad);

Status DynamicStitchGrad(const Scope& scope, const Operation& op,
                         const std::vector<Output>& grad_inputs,
                         std::vector<Output>* grad_outputs) {
  // Running example:
  // indices = {2, [1, 0]}
  // data = {[d_1, d_2], [[d_3, d_4], [d_5, d_6]]}
  // out = [[d_5, d_6], [d_3, d_4], [d_1, d_2]]
  // grad = [[g_1, g_2], [g_3, g_4], [g_5, g_6]]

  // indices and data are two equal-sized lists passed
  // into DynamicStitch.
  // num_values = 2
  int32 num_values = op.num_inputs() / 2;

  // Stop propagation along the indices list
  for (int32 i = 0; i < num_values; i++) {
    grad_outputs->push_back(NoGradient());
  }

  // DynamicStitch shuffles its data to the output (using items in
  // indices) so the gradient propagated to a given data input simply
  // selects the gradient for its output position.
  for (int32 i = 0; i < num_values; i++) {
    // index has the destination positions for the i'th data
    // element. We cast it into an int32 if necessary, so we can use
    // it from a Gather op.
    // i = 0: index = 2
    // i = 1: index = [1, 0]
    auto index = op.input(i);
    if (index.type() != DT_INT32) {
      index = Cast(scope, index, DT_INT32);
    }
    // Gather the index specified locations in the gradient and
    // propagate it as the gradient for the i'th data item.
    // i = 0: gather(grad, 2) = [g_5, g_6]
    // i = 1: gather(grad, [1, 0]) = [[g_3, g_4], [g_1, g_2]]
    grad_outputs->push_back(Gather(scope, grad_inputs[0], index));
  }

  return scope.status();
}
REGISTER_GRADIENT_OP("DynamicStitch", DynamicStitchGrad);
REGISTER_GRADIENT_OP("DynamicStitchFast", DynamicStitchGrad);
REGISTER_GRADIENT_OP("ParallelDynamicStitch", DynamicStitchGrad);