|
| 1 | +/* Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved. |
| 2 | +
|
| 3 | +Licensed under the Apache License, Version 2.0 (the "License"); |
| 4 | +you may not use this file except in compliance with the License. |
| 5 | +You may obtain a copy of the License at |
| 6 | +
|
| 7 | + http://www.apache.org/licenses/LICENSE-2.0 |
| 8 | +
|
| 9 | +Unless required by applicable law or agreed to in writing, software |
| 10 | +distributed under the License is distributed on an "AS IS" BASIS, |
| 11 | +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 12 | +See the License for the specific language governing permissions and |
| 13 | +limitations under the License. */ |
| 14 | + |
| 15 | +#include "paddle/fluid/operators/trunc_op.h" |
| 16 | + |
| 17 | +namespace paddle { |
| 18 | +namespace operators { |
| 19 | + |
| 20 | +class TruncOp : public framework::OperatorWithKernel { |
| 21 | + public: |
| 22 | + using framework::OperatorWithKernel::OperatorWithKernel; |
| 23 | + |
| 24 | + void InferShape(framework::InferShapeContext *ctx) const override { |
| 25 | + OP_INOUT_CHECK(ctx->HasInput("X"), "Input", "X", "trunc"); |
| 26 | + OP_INOUT_CHECK(ctx->HasOutput("Out"), "Output", "Out", "trunc"); |
| 27 | + auto input_dims = ctx->GetInputDim("X"); |
| 28 | + ctx->SetOutputDim("Out", input_dims); |
| 29 | + ctx->ShareLoD("X", /*->*/ "Out"); |
| 30 | + } |
| 31 | +}; |
| 32 | + |
| 33 | +class TruncOpMaker : public framework::OpProtoAndCheckerMaker { |
| 34 | + public: |
| 35 | + void Make() override { |
| 36 | + AddInput("X", "(Tensor), The input tensor of trunc op."); |
| 37 | + AddOutput("Out", "(Tensor), The output tensor of trunc op."); |
| 38 | + AddComment(R"DOC( |
| 39 | +Trunc Operator. |
| 40 | +Returns a new tensor with the truncated integer values of input. |
| 41 | +$$out = trunc(x)$$ |
| 42 | +)DOC"); |
| 43 | + } |
| 44 | +}; |
| 45 | + |
| 46 | +class TruncGradOp : public framework::OperatorWithKernel { |
| 47 | + public: |
| 48 | + using framework::OperatorWithKernel::OperatorWithKernel; |
| 49 | + |
| 50 | + void InferShape(framework::InferShapeContext *ctx) const override { |
| 51 | + OP_INOUT_CHECK(ctx->HasInput(framework::GradVarName("Out")), "Input", |
| 52 | + framework::GradVarName("Out"), "TruncGrad"); |
| 53 | + OP_INOUT_CHECK(ctx->HasOutput(framework::GradVarName("X")), "Output", |
| 54 | + framework::GradVarName("X"), "TruncGrad"); |
| 55 | + |
| 56 | + auto dout_dims = ctx->GetInputDim(framework::GradVarName("Out")); |
| 57 | + ctx->SetOutputDim(framework::GradVarName("X"), dout_dims); |
| 58 | + } |
| 59 | +}; |
| 60 | + |
| 61 | +template <typename T> |
| 62 | +class TruncGradOpMaker : public framework::SingleGradOpMaker<T> { |
| 63 | + public: |
| 64 | + using framework::SingleGradOpMaker<T>::SingleGradOpMaker; |
| 65 | + |
| 66 | + void Apply(GradOpPtr<T> retv) const override { |
| 67 | + retv->SetType("trunc_grad"); |
| 68 | + retv->SetInput(framework::GradVarName("Out"), this->OutputGrad("Out")); |
| 69 | + retv->SetAttrMap(this->Attrs()); |
| 70 | + retv->SetOutput(framework::GradVarName("X"), this->InputGrad("X")); |
| 71 | + } |
| 72 | +}; |
| 73 | + |
| 74 | +} // namespace operators |
| 75 | +} // namespace paddle |
| 76 | + |
| 77 | +namespace ops = paddle::operators; |
| 78 | +REGISTER_OPERATOR(trunc, ops::TruncOp, ops::TruncOpMaker, |
| 79 | + ops::TruncGradOpMaker<paddle::framework::OpDesc>, |
| 80 | + ops::TruncGradOpMaker<paddle::imperative::OpBase>); |
| 81 | + |
| 82 | +REGISTER_OPERATOR(trunc_grad, ops::TruncGradOp); |
| 83 | + |
| 84 | +REGISTER_OP_CPU_KERNEL(trunc, ops::TruncKernel<float>, ops::TruncKernel<double>, |
| 85 | + ops::TruncKernel<int>, ops::TruncKernel<int64_t>); |
| 86 | + |
| 87 | +REGISTER_OP_CPU_KERNEL(trunc_grad, ops::TruncGradKernel<float>, |
| 88 | + ops::TruncGradKernel<double>, ops::TruncGradKernel<int>, |
| 89 | + ops::TruncGradKernel<int64_t>); |
0 commit comments