Add gradient for tensorflow::ops::Fill
See https://github.com/tensorflow/tensorflow/issues/20926
This commit is contained in:
parent
9b3a29889f
commit
b462815ccc
@ -120,6 +120,24 @@ Status SplitGrad(const Scope& scope, const Operation& op,
|
|||||||
}
|
}
|
||||||
REGISTER_GRADIENT_OP("Split", SplitGrad);
|
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,
|
Status DiagGrad(const Scope& scope, const Operation& op,
|
||||||
const std::vector<Output>& grad_inputs,
|
const std::vector<Output>& grad_inputs,
|
||||||
std::vector<Output>* grad_outputs) {
|
std::vector<Output>* grad_outputs) {
|
||||||
|
@ -108,6 +108,14 @@ TEST_F(ArrayGradTest, SplitGrad) {
|
|||||||
RunTest({x}, {x_shape}, y.output, {y_shape, y_shape});
|
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) {
|
TEST_F(ArrayGradTest, DiagGrad) {
|
||||||
TensorShape x_shape({5, 2});
|
TensorShape x_shape({5, 2});
|
||||||
auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(x_shape));
|
auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(x_shape));
|
||||||
|
Loading…
Reference in New Issue
Block a user