Lines Matching defs:ms
263 xla::ComputationDataHandle var, ms, mom;
265 OP_REQUIRES_OK(ctx, ctx->ReadVariableInput(1, type, &ms_shape, &ms));
286 // var should be the same shape as mom and ms.
288 errors::InvalidArgument("var and ms do not have the same shape",
306 // ms <- rho * ms_{t-1} + (1-rho) * grad * grad
307 // mom <- momentum * mom_{t-1} + lr * grad / sqrt(ms + epsilon)
310 // We use an alternate formulation of the ms equation:
312 // ms <- ms + (grad**2 - ms) * (1 - rho)
316 // ms <- ms + grad**2 - rho * grad ** 2 - ms + ms * rho
320 // ms <- grad**2 (1 - rho) + ms * rho
324 ms,
325 b->Mul(b->Sub(b->Pow(grad, XlaHelpers::FloatLiteral(b, type, 2.0)), ms),
