Caffe2 - C++ API
A deep learning, cross platform ML framework
batch_matmul.cc
1 #include "caffe2/operators/experimental/c10/schemas/batch_matmul.h"
2 #include <ATen/core/dispatch/OpSchemaRegistration.h>
3 #include "caffe2/core/operator_c10wrapper.h"
4 
6 
7 namespace caffe2 {
8 namespace ops {
9 // TODO Parse schema string instead of creating FunctionSchema manually
10 C10_DEFINE_OP_SCHEMA(
11  BatchMatmul,
12  FunctionSchema(
13  "_c10_experimental::BatchMatmul",
14  "",
15  (std::vector<c10::Argument>{
16  c10::Argument("A"),
17  c10::Argument("B"),
18  c10::Argument("output"),
19  c10::Argument("trans_a", IntType::get()),
20  c10::Argument("trans_b", IntType::get()),
21  c10::Argument("broadcast", IntType::get())}),
22  (std::vector<c10::Argument>{})));
23 }
24 }
25 
26 namespace caffe2 {
27 
28 REGISTER_C10_OPERATOR_FOR_CAFFE2_DISPATCH_CPU(
29  ops::BatchMatmul(),
30  C10BatchMatMul_DontUseThisOpYet)
31 }
The CPU Context, representing the bare minimum of what a Context class in Caffe2 should implement...
Definition: context.h:40
A global dictionary that holds information about what Caffe2 modules have been loaded in the current ...
Definition: blob.h:13