forked from pytorch/pytorch
-
Notifications
You must be signed in to change notification settings - Fork 0
/
alias_with_name.cc
27 lines (22 loc) · 980 Bytes
/
alias_with_name.cc
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
#include "caffe2/operators/alias_with_name.h"
namespace caffe2 {
// NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables)
REGISTER_CPU_OPERATOR(AliasWithName, AliasWithNameOp<CPUContext>);
// NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables)
OPERATOR_SCHEMA(AliasWithName)
.NumInputs(1)
.NumOutputs(1)
.AllowInplace({{0, 0}})
.IdenticalTypeAndShape()
.SetDoc(R"DOC(
Similar with AliasOp, storing the alias name as operator argument.
)DOC")
.Arg("name", "name of the aliasing")
.Arg("is_backward", "weather or not to alias forward or backward")
.Input(0, "input", "Input tensor whose storage will be shared.")
.Output(0, "output", "Tensor of same shape as input, sharing its storage.");
} // namespace caffe2
C10_EXPORT_CAFFE2_OP_TO_C10_CPU(
AliasWithName,
"_caffe2::AliasWithName(Tensor input, str name, bool is_backward = False) -> (Tensor output)",
caffe2::AliasWithNameOp<caffe2::CPUContext>);