| 92 | } |
| 93 | |
| 94 | Maybe<void> InferSGDUpdateTensorDesc(user_op::InferContext* ctx) { |
| 95 | const user_op::TensorDesc& model = ctx->InputTensorDesc("model", 0); |
| 96 | const Shape& shape = model.shape(); |
| 97 | const user_op::TensorDesc& model_diff = ctx->InputTensorDesc("model_diff", 0); |
| 98 | if (shape.NumAxes() > 0 && model_diff.shape().NumAxes() > 0) { |
| 99 | CHECK_EQ_OR_RETURN(model_diff.shape(), shape); |
| 100 | } |
| 101 | JUST(CheckLearningRateShape(ctx)); |
| 102 | if (ctx->has_input("model_copy", 0)) { |
| 103 | CHECK_EQ_OR_RETURN(ctx->InputTensorDesc("model_copy", 0).shape(), shape) |
| 104 | << "Model copy shape should be equal to Model shape. "; |
| 105 | } |
| 106 | if (ctx->has_input("scale_by_tensor", 0)) { |
| 107 | const auto& scale_by_tensor = ctx->InputTensorDesc("scale_by_tensor", 0); |
| 108 | JUST(CheckScalarShape(&scale_by_tensor)); |
| 109 | } |
| 110 | return Maybe<void>::Ok(); |
| 111 | } |
| 112 | Maybe<void> InferSGDUpdateDataType(user_op::InferContext* ctx) { |
| 113 | JUST(CheckLearningRateDataType(ctx)); |
| 114 | if (ctx->has_input("scale_by_tensor", 0)) { |
no test coverage detected