Function InferRmsPropUpdateTensorDesc

oneflow/user/ops/model_update_ops.cpp:440–461

}

Maybe<void> InferRmsPropUpdateTensorDesc(user_op::InferContext* ctx) {
  const user_op::TensorDesc& model = ctx->InputTensorDesc("model", 0);

  const Shape& shape = model.shape();
  const user_op::TensorDesc& model_diff = ctx->InputTensorDesc("model_diff", 0);
  CHECK_EQ_OR_RETURN(model_diff.shape(), shape);
  const user_op::TensorDesc& mean_square = ctx->InputTensorDesc("mean_square", 0);
  JUST(CheckShapeLike(&mean_square, &model));
  JUST(CheckLearningRateShape(ctx));
  if (ctx->has_input("scale_by_tensor", 0)) {
    const auto& scale_by_tensor = ctx->InputTensorDesc("scale_by_tensor", 0);
    JUST(CheckScalarShape(&scale_by_tensor));
  }
  if (ctx->Attr<bool>("centered")) {
    CHECK_OR_RETURN(ctx->has_input("mean_gradient", 0));
    const user_op::TensorDesc& mean_gradient = ctx->InputTensorDesc("mean_gradient", 0);
    JUST(CheckShapeLike(&mean_gradient, &model));
  } else {
    CHECK_OR_RETURN(!ctx->has_input("mean_gradient", 0));
  }
  return Maybe<void>::Ok();
}

Maybe<void> InferRmsPropUpdateDataType(user_op::InferContext* ctx) {
  const user_op::TensorDesc& model = ctx->InputTensorDesc("model", 0);