Add gradient for tensorflow::ops::Fill

See https://github.com/tensorflow/tensorflow/issues/20926
This commit is contained in:
KB Sriram 2018-07-18 08:31:46 -07:00
parent 9b3a29889f
commit b462815ccc
2 changed files with 26 additions and 0 deletions

View File

@ -120,6 +120,24 @@ Status SplitGrad(const Scope& scope, const Operation& op,
}
REGISTER_GRADIENT_OP("Split", SplitGrad);
Status FillGrad(const Scope& scope, const Operation& op,
const std::vector<Output>& grad_inputs,
std::vector<Output>* grad_outputs) {
// y = fill(fill_shape, x)
// No gradient returned for the fill_shape argument.
grad_outputs->push_back(NoGradient());
// The gradient for x (which must be a scalar) is just the sum of
// all the gradients from the shape it fills.
// We use ReduceSum to implement this, which needs an argument providing
// the indices of all the dimensions of the incoming gradient.
// grad(x) = reduce_sum(grad(y), [0..rank(grad(y))])
auto all_dims = Range(scope, Const(scope, 0), Rank(scope, grad_inputs[0]),
Const(scope, 1));
grad_outputs->push_back(ReduceSum(scope, grad_inputs[0], all_dims));
return scope.status();
}
REGISTER_GRADIENT_OP("Fill", FillGrad);
Status DiagGrad(const Scope& scope, const Operation& op,
const std::vector<Output>& grad_inputs,
std::vector<Output>* grad_outputs) {

View File

@ -108,6 +108,14 @@ TEST_F(ArrayGradTest, SplitGrad) {
RunTest({x}, {x_shape}, y.output, {y_shape, y_shape});
}
TEST_F(ArrayGradTest, FillGrad) {
TensorShape x_shape({});
auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(x_shape));
TensorShape y_shape({2, 5, 3});
auto y = Fill(scope_, {2, 5, 3}, x);
RunTest(x, x_shape, y, y_shape);
}
TEST_F(ArrayGradTest, DiagGrad) {
TensorShape x_shape({5, 2});
auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(x_shape));