[TF:XLA] Use a more numerically accurate formulation of ResourceApplyRMSProp.

PiperOrigin-RevId: 161706120
This commit is contained in:
Peter Hawkins 2017-07-12 13:10:45 -07:00 committed by TensorFlower Gardener
parent 45a58d378e
commit e424d209a7

View File

@ -352,9 +352,9 @@ class ResourceApplyRMSProp : public XlaOpKernel {
b->Sub(XlaHelpers::FloatLiteral(b, type, 1.0), rho))); b->Sub(XlaHelpers::FloatLiteral(b, type, 1.0), rho)));
xla::ComputationDataHandle new_mom = xla::ComputationDataHandle new_mom =
b->Add(b->Mul(mom, momentum), b->Add(b->Mul(mom, momentum),
b->Div(b->Mul(grad, lr), b->Mul(b->Mul(grad, lr),
b->Pow(b->Add(new_ms, epsilon), b->Pow(b->Add(new_ms, epsilon),
XlaHelpers::FloatLiteral(b, type, 0.5)))); XlaHelpers::FloatLiteral(b, type, -0.5))));
xla::ComputationDataHandle new_var = b->Sub(var, new_mom); xla::ComputationDataHandle new_var = b->Sub(var, new_mom);
OP_REQUIRES_OK(ctx, ctx->AssignVariable(0, type, new_var)); OP_REQUIRES_OK(ctx, ctx->AssignVariable(0, type, new_var));