aboutsummaryrefslogtreecommitdiffhomepage
path: root/tensorflow/cc/gradients/math_grad.cc
diff options
context:
space:
mode:
Diffstat (limited to 'tensorflow/cc/gradients/math_grad.cc')
-rw-r--r--tensorflow/cc/gradients/math_grad.cc20
1 files changed, 20 insertions, 0 deletions
diff --git a/tensorflow/cc/gradients/math_grad.cc b/tensorflow/cc/gradients/math_grad.cc
index cd215f740d..1329b568ab 100644
--- a/tensorflow/cc/gradients/math_grad.cc
+++ b/tensorflow/cc/gradients/math_grad.cc
@@ -1022,6 +1022,26 @@ Status ProdGrad(const Scope& scope, const Operation& op,
}
REGISTER_GRADIENT_OP("Prod", ProdGrad);
+Status SegmentSumGrad(const Scope& scope, const Operation& op,
+ const std::vector<Output>& grad_inputs,
+ std::vector<Output>* grad_outputs) {
+ // The SegmentSum operation sums segments of the Tensor that have the same
+ // index in the segment_ids parameter.
+ // i.e z = [2, 3, 4, 5], segment_ids [0, 0, 0, 1]
+ // will produce [2 + 3 + 4, 5] = [9, 5]
+ // The gradient that will flow back to the gather operation will look like
+ // [x1, x2], it will have the same shape as the output of the SegmentSum
+ // operation. The differentiation step of the SegmentSum operation just
+ // broadcast the gradient in order to retrieve the z's shape.
+ // dy/dz = [x1, x1, x1, x2]
+ grad_outputs->push_back(Gather(scope, grad_inputs[0], op.input(1)));
+
+ // stop propagation along segment_ids
+ grad_outputs->push_back(NoGradient());
+ return scope.status();
+}
+REGISTER_GRADIENT_OP("SegmentSum", SegmentSumGrad);
+
// MatMulGrad helper function used to compute two MatMul operations
// based on input matrix transposition combinations.
Status MatMulGradHelper(const Scope& scope, const bool is_batch,