Function InferAdamUpdateTensorDesc

oneflow/user/ops/model_update_ops.cpp:185–204

}

Maybe<void> InferAdamUpdateTensorDesc(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& m = ctx->InputTensorDesc("m", 0);
  JUST(CheckShapeLike(&m, &model));
  const user_op::TensorDesc& v = ctx->InputTensorDesc("v", 0);
  JUST(CheckShapeLike(&v, &model));
  JUST(CheckLearningRateShape(ctx));
  if (ctx->has_input("model_copy", 0)) {
    CHECK_EQ_OR_RETURN(ctx->InputTensorDesc("model_copy", 0).shape(), shape)
        << "Model copy shape should be equal to Model shape. ";
  }
  if (ctx->has_input("scale_by_tensor", 0)) {
    const auto& scale_by_tensor = ctx->InputTensorDesc("scale_by_tensor", 0);
    JUST(CheckScalarShape(&scale_by_tensor));
  }
  return Maybe<void>::Ok();
}
Maybe<void> InferAdamUpdateDataType(user_op::InferContext* ctx) {
  const user_op::TensorDesc& model = ctx->InputTensorDesc("model", 0);
  const user_op::TensorDesc& m = ctx->InputTensorDesc("m", 0);