pytorch  1.8.2 About: PyTorch provides Tensor computation (like NumPy) with strong GPU acceleration and Deep Neural Networks (in Python) built on a tape-based autograd system. LTS (Long Term Support) release.   Fossies Dox: pytorch-1.8.2.tar.gz  ("unofficial" and yet experimental doxygen-generated source code documentation)
matmul_op.cc
Go to the documentation of this file.
2
3namespace caffe2 {
4
6
7OPERATOR_SCHEMA(MatMul)
8 .NumInputs(2, 3)
9 .NumOutputs(1)
10 .TensorInferenceFunction([](const OperatorDef& def,
11 const vector<TensorShape>& in) {
13 out[0].set_data_type(in[0].data_type());
15 int axis_a = arg_helper.GetSingleArgument<int>("axis_a", 1);
16 int axis_b = arg_helper.GetSingleArgument<int>("axis_b", 1);
17 int trans_a = arg_helper.GetSingleArgument<bool>("trans_a", false);
18 int trans_b = arg_helper.GetSingleArgument<bool>("trans_b", false);
21
24 if (trans_a) {
26 }
27 if (trans_b) {
29 }
30
33
34 return out;
35 })
36 .SetDoc(R"DOC(
37Matrix multiplication \$Y = A * B\$, where `A` has size (M x K), `B` has size
38(K x N), and `Y` will have a size (M x N). To transpose `A` or `B` before
39multiplication, pass 1 to the `trans_a` and/or `trans_b` arguments, which
40separate the first and second dimensions of the respective matrices using
41`axis_a` and `axis_b`.
42
44
45- https://github.com/pytorch/pytorch/blob/master/caffe2/operators/matmul_op.cc
46
47<details>
48
49<summary> <b>Example</b> </summary>
50
51**Code**
52
53```
54workspace.ResetWorkspace()
55
56op = core.CreateOperator(
57 "MatMul",
58 ["A", "B"],
59 ["Y"],
60)
61
62workspace.FeedBlob("A", np.random.randint(10, size=(3,3)).astype(np.float32))
63workspace.FeedBlob("B", np.random.randint(10, size=(3,3)).astype(np.float32))
64print("A:", workspace.FetchBlob("A"))
65print("B:", workspace.FetchBlob("B"))
66workspace.RunOperatorOnce(op)
67print("Y:", workspace.FetchBlob("Y"))
68```
69
70**Result**
71
72```
73A: [[1. 8. 3.]
74 [6. 4. 4.]
75 [5. 4. 7.]]
76B: [[4. 0. 3.]
77 [3. 1. 1.]
78 [8. 5. 8.]]
79Y: [[52. 23. 35.]
80 [68. 24. 54.]
81 [88. 39. 75.]]
82```
83
84</details>
85
86)DOC")
87 .Input(
88 0,
89 "A",
90 "*(type: Tensor`<float>`)* 2D matrix of size (M x K).")
91 .Input(
92 1,
93 "B",
94 "*(type: Tensor`<float>`)* 2D matrix of size (K x N).")
95 .Output(
96 0,
97 "Y",
98 "*(type: Tensor`<float>`)* 2D matrix of size (M x N).")
99 .Arg(
100 "axis_a",
101 "*(type: int; default: 1)* Exclusive axis that divides the first and "
102 "second dimension of matrix `A`.")
103 .Arg(
104 "axis_b",
105 "*(type: int; default: 1)* Exclusive axis that divides the first and "
106 "second dimension of matrix `B`.")
107 .Arg(
108 "trans_a",
109 "*(type: int; default: 0)* Pass 1 to transpose `A` before multiplication and "
110 "after the dimension adjustment using `axis_a`.")
111 .Arg(
112 "trans_b",
113 "*(type: int; default: 0)* Pass 1 to transpose `B` before multiplication and "
114 "after the dimension adjustment using `axis_b`.");
115
119 CAFFE_ENFORCE(def_.input_size() == 2 || def_.input_size() == 3);
120
121 bool axis_a = 1;
122 bool axis_b = 1;
123 bool trans_a = 0;
124 bool trans_b = 0;
125
126 if (ArgumentHelper::HasArgument(Def(), "trans_a")) {
127 trans_a = GetArgument(Def(), "trans_a").i();
128 }
129 if (ArgumentHelper::HasArgument(Def(), "trans_b")) {
130 trans_b = GetArgument(Def(), "trans_b").i();
131 }
132 if (ArgumentHelper::HasArgument(Def(), "axis_a")) {
133 axis_a = GetArgument(Def(), "axis_a").i();
134 }
135 if (ArgumentHelper::HasArgument(Def(), "axis_b")) {
136 axis_b = GetArgument(Def(), "axis_b").i();
137 }
138
139 if (trans_a) {
140 if (trans_b) {
141 // A'B':
142 // dA = B'G', dB = G'A'
143 return vector<OperatorDef>{
145 "MatMul",
146 "",
147 vector<string>{I(1), GO(0), I(0)},
148 vector<string>{GI(0)},
149 vector<Argument>{MakeArgument<int>("trans_a", 1),
150 MakeArgument<int>("trans_b", 1),
151 MakeArgument<int>("axis_a", axis_b)}),
153 "MatMul",
154 "",
155 vector<string>{GO(0), I(0), I(1)},
156 vector<string>{GI(1)},
157 vector<Argument>{MakeArgument<int>("trans_a", 1),
158 MakeArgument<int>("trans_b", 1),
159 MakeArgument<int>("axis_b", axis_a)})};
160 } else {
161 // A'B:
162 // dA = BG', dB = AG
163 return vector<OperatorDef>{
165 "MatMul",
166 "",
167 vector<string>{I(1), GO(0), I(0)},
168 vector<string>{GI(0)},
169 vector<Argument>{MakeArgument<int>("trans_b", 1),
170 MakeArgument<int>("axis_a", axis_b)}),
172 "MatMul",
173 "",
174 vector<string>{I(0), GO(0), I(1)},
175 vector<string>{GI(1)},
176 vector<Argument>{MakeArgument<int>("axis_a", axis_a)})};
177 }
178 } else {
179 if (trans_b) {
180 // AB':
181 // dA = GB, dB = G'A
182 return vector<OperatorDef>{
184 "MatMul",
185 "",
186 vector<string>{GO(0), I(1), I(0)},
187 vector<string>{GI(0)},
188 vector<Argument>{MakeArgument<int>("axis_b", axis_b)}),
190 "MatMul",
191 "",
192 vector<string>{GO(0), I(0), I(1)},
193 vector<string>{GI(1)},
194 vector<Argument>{MakeArgument<int>("trans_a", 1),
195 MakeArgument<int>("axis_b", axis_a)})};
196 } else {
197 // AB:
198 // dA = GB', dB = A'G
199 return vector<OperatorDef>{
201 "MatMul",
202 "",
203 vector<string>{GO(0), I(1), I(0)},
204 vector<string>{GI(0)},
205 vector<Argument>{MakeArgument<int>("trans_b", 1),
206 MakeArgument<int>("axis_b", axis_b)}),
208 "MatMul",
209 "",
210 vector<string>{I(0), GO(0), I(1)},
211 vector<string>{GI(1)},
212 vector<Argument>{MakeArgument<int>("trans_a", 1),
213 MakeArgument<int>("axis_a", axis_a)})};
214 }
215 }
216 }
217
218 bool CopyArguments() const override {
219 return false;
220 }
221};
222
224
225} // namespace caffe2
A helper class to index into arguments.
Definition: proto_utils.h:203
static bool HasArgument(const Def &def, const string &name)
Definition: proto_utils.h:206
static T GetSingleArgument(const Def &def, const string &name, const T &default_value)
Definition: proto_utils.h:211
bool CopyArguments() const override
Definition: matmul_op.cc:168
Definition: matmul_op.cc:68
const OperatorDef & Def() const
const OperatorDef & def_
int canonical_axis_index_(int axis_index, int ndims)
Definition: TensorImpl.h:93
int64_t size_to_dim_(int k, IntArrayRef dims)
Definition: TensorImpl.h:67
int64_t size_from_dim_(int k, IntArrayRef dims)
Return product of all dimensions starting from k.
Definition: TensorImpl.h:58
Definition: blob.h:13
REGISTER_CPU_OPERATOR(ATen, ATenOp< CPUContext >)
C10_EXPORT const Argument & GetArgument(const OperatorDef &def, const string &name)
Definition: proto_utils.cc:522
OPERATOR_SCHEMA(ATen)
int canonical_axis_b
Definition: matmul_op.cc:20
return vector< TensorShape >
Definition: slice_op.cc:110
we first initialize the output tensor to all and then do accumulation Any further calls to the The input tensor that has to be accumulated to the output tensor If the output size is not the same as input size
const vector< TensorShape > & in
ArgumentHelper arg_helper(def)
int trans_b
Definition: matmul_op.cc:18
vector< int64_t > GetDimsVector(const TensorShape &shape)
int canonical_axis_a
Definition: matmul_op.cc:19
int axis_b
Definition: matmul_op.cc:16
*and produces a single output tensor *expanded *The op also takes an argument *dims *with a list of dimensions for where to add the single dimensional entries If the same blob is provided as input and the operation is copy free This is the exact inverse operation of *Squeeze *Github dims
int M
Definition: matmul_op.cc:22
SparseLengths8BitsRowwiseOp< CPUContext, 0, 1 >::LENGTHS SetDoc(R"DOC( Variation of SparseLengthsMean operator, where DATA is stored using 8bits. DATA was quantized with 8Bit row-wise quantization (see doc to FloatToRowwiseQuantized8Bits operator). To restore DATA from 8Bit, we use additional input that stores scales and biases. )DOC") .Input(0