assign_value_op.cc 2.6 KB
Newer Older
1
//   Copyright (c) 2018 PaddlePaddle Authors. All Rights Reserved.
D
dzhwinter 已提交
2 3 4 5 6 7 8 9 10 11 12 13
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
X
xuwei06 已提交
14

Y
Yi Wang 已提交
15
#include "paddle/fluid/operators/assign_value_op.h"
16 17
#include <string>
#include <vector>
X
xuwei06 已提交
18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37

namespace paddle {
namespace operators {

class AssignValueOp : public framework::OperatorWithKernel {
 public:
  AssignValueOp(const std::string &type,
                const framework::VariableNameMap &inputs,
                const framework::VariableNameMap &outputs,
                const framework::AttributeMap &attrs)
      : OperatorWithKernel(type, inputs, outputs, attrs) {}

  void InferShape(framework::InferShapeContext *ctx) const override {
    PADDLE_ENFORCE(ctx->HasOutput("Out"),
                   "Output(Out) of AssignValueOp should not be null.");
    auto shape = ctx->Attrs().Get<std::vector<int>>("shape");
    ctx->SetOutputDim("Out", framework::make_ddim(shape));
  }

 protected:
38
  framework::OpKernelType GetExpectedKernelType(
X
xuwei06 已提交
39 40
      const framework::ExecutionContext &ctx) const override {
    return framework::OpKernelType(
41 42
        framework::proto::VarType::Type(ctx.Attr<int>("dtype")),
        ctx.GetPlace());
X
xuwei06 已提交
43 44 45 46 47
  }
};

class AssignValueOpMaker : public framework::OpProtoAndCheckerMaker {
 public:
Y
Yu Yang 已提交
48
  void Make() override {
X
xuwei06 已提交
49 50 51 52 53
    AddOutput("Out", "(Tensor) Output tensor of assign_value operator.");
    AddAttr<std::vector<int>>("shape",
                              "(vector<int>) "
                              "Shape of values.");
    AddAttr<int>("dtype", "data type of values")
54 55
        .InEnum({framework::proto::VarType::INT32,
                 framework::proto::VarType::FP32});
X
xuwei06 已提交
56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73
    AddAttr<std::vector<float>>("fp32_values", "store the float values")
        .SetDefault({});
    AddAttr<std::vector<int>>("int32_values", "store the int values")
        .SetDefault({});
    AddComment(R"DOC(
AssignValue operator

$$Out = values$$
)DOC");
  }
};

}  // namespace operators
}  // namespace paddle

namespace ops = paddle::operators;

REGISTER_OPERATOR(assign_value, ops::AssignValueOp, ops::AssignValueOpMaker);
74 75
REGISTER_OP_CPU_KERNEL(assign_value, ops::AssignValueKernel<int>,
                       ops::AssignValueKernel<float>);
新手
引导
客服 返回
顶部