Function UnsortedSegmentMinOrMaxGrad

tensorflow/cc/gradients/math_grad.cc:1288–1320

}

Status UnsortedSegmentMinOrMaxGrad(const Scope& scope, const Operation& op,
                                   const std::vector<Output>& grad_inputs,
                                   std::vector<Output>* grad_outputs) {
  if (op.num_inputs() != 3) {
    return errors::InvalidArgument("UnsortedSegmentMax requires 3 arguments");
  }

  if (grad_inputs.size() != 1) {
    return errors::InvalidArgument(
        "UnsortedSegmentMax grad requires 1 grad input");
  }

  auto grad = grad_inputs[0];
  // Get the number of selected (minimum or maximum) elements in each segment.
  auto zero_clipped_indices = GetZeroClippedIndices(scope, op.input(1));
  auto is_positive = GetIsPositive(scope, op.output(0), op.input(1));
  Output gathered_outputs = GatherDropNegatives(
      scope, op.output(0), zero_clipped_indices, is_positive);
  Output is_selected = Equal(scope, op.input(0), gathered_outputs);
  is_selected = LogicalAnd(scope, is_selected, is_positive);
  auto num_selected = UnsortedSegmentSum(
      scope, Cast(scope, is_selected, grad.type()), op.input(1), op.input(2));
  // Compute the gradient for each segment.The gradient for the ith segment is
  // divided evenly among the selected elements in that segment.
  auto weighted_grads = Div(scope, grad, num_selected);
  auto gathered_grads = GatherDropNegatives(scope, weighted_grads,
                                            zero_clipped_indices, is_positive);
  auto zeros = ZerosLike(scope, gathered_grads);
  grad_outputs->push_back(SelectV2(scope, is_selected, gathered_grads, zeros));
  grad_outputs->push_back(NoGradient());
  grad_outputs->push_back(NoGradient());
  return scope.status();
}

REGISTER_GRADIENT_OP("UnsortedSegmentMax", UnsortedSegmentMinOrMaxGrad);
REGISTER_GRADIENT_OP("UnsortedSegmentMin", UnsortedSegmentMinOrMaxGrad);