Infer schema for experimental ops (pytorch#20513)

Summary: Pull Request resolved: pytorch#20513 They've been using an old API, switch them to the new one instead. Reviewed By: li-roy Differential Revision: D15346349 fbshipit-source-id: 538eb460897ec6addebeebf88b316eb0d6b1dd6f
IvanKobzarev · May 16, 2019 · 09f22d1 · 09f22d1
1 parent 9bd3305
commit 09f22d1
Show file tree

Hide file tree

Showing 17 changed files with 21 additions and 188 deletions.
diff --git a/caffe2/operators/experimental/c10/cpu/add_cpu.cc b/caffe2/operators/experimental/c10/cpu/add_cpu.cc
@@ -71,16 +71,7 @@ void add_op_cpu_impl(
 }
 
 static auto registry = c10::RegisterOperators().op(
-    FunctionSchema(
-        "_c10_experimental::Add",
-        "",
-        (std::vector<c10::Argument>{
-            c10::Argument("input1"),
-            c10::Argument("input2"),
-            c10::Argument("output"),
-            c10::Argument("legacy_broadcast", BoolType::get()),
-            c10::Argument("axis", IntType::get())}),
-        (std::vector<c10::Argument>{})),
+    "_c10_experimental::Add",
     c10::kernel<decltype(add_op_cpu_impl<float>), &add_op_cpu_impl<float>>(),
     c10::dispatchKey(CPUTensorId()));
 

diff --git a/caffe2/operators/experimental/c10/cpu/averaged_loss_cpu.cc b/caffe2/operators/experimental/c10/cpu/averaged_loss_cpu.cc
@@ -44,12 +44,7 @@ class averaged_loss_cpu final : public c10::OperatorKernel {
 };
 
 static auto registry = c10::RegisterOperators().op(
-    FunctionSchema(
-        "_c10_experimental::AveragedLoss",
-        "",
-        (std::vector<c10::Argument>{c10::Argument("input"),
-                                    c10::Argument("output")}),
-        (std::vector<c10::Argument>{})),
+    "_c10_experimental::AveragedLoss",
     c10::kernel<averaged_loss_cpu<float, CPUContext>>(),
     c10::dispatchKey(CPUTensorId()));
 

diff --git a/caffe2/operators/experimental/c10/cpu/batch_gather_cpu.cc b/caffe2/operators/experimental/c10/cpu/batch_gather_cpu.cc
@@ -65,13 +65,7 @@ void batch_gather_op_cpu(const at::Tensor& data,
 }
 
 static auto registry = c10::RegisterOperators().op(
-    FunctionSchema(
-        "_c10_experimental::BatchGather",
-        "",
-        (std::vector<c10::Argument>{c10::Argument("data"),
-                                    c10::Argument("indices"),
-                                    c10::Argument("output")}),
-        (std::vector<c10::Argument>{})),
+    "_c10_experimental::BatchGather",
     c10::kernel<decltype(batch_gather_op_cpu), &batch_gather_op_cpu>(),
     c10::dispatchKey(CPUTensorId()));
 

diff --git a/caffe2/operators/experimental/c10/cpu/batch_matmul_cpu.cc b/caffe2/operators/experimental/c10/cpu/batch_matmul_cpu.cc
@@ -269,17 +269,7 @@ class batch_matmul_cpu final : public c10::OperatorKernel {
 };
 
 static auto registry = c10::RegisterOperators().op(
-    FunctionSchema(
-        "_c10_experimental::BatchMatmul",
-        "",
-        (std::vector<c10::Argument>{
-            c10::Argument("A"),
-            c10::Argument("B"),
-            c10::Argument("output"),
-            c10::Argument("trans_a", IntType::get()),
-            c10::Argument("trans_b", IntType::get()),
-            c10::Argument("broadcast", IntType::get())}),
-        (std::vector<c10::Argument>{})),
+    "_c10_experimental::BatchMatmul",
     c10::kernel<batch_matmul_cpu<float, CPUContext>>(),
     c10::dispatchKey(CPUTensorId()));
 

diff --git a/caffe2/operators/experimental/c10/cpu/cast_cpu.cc b/caffe2/operators/experimental/c10/cpu/cast_cpu.cc
@@ -87,15 +87,7 @@ void cast_op_cpu(
 }
 
 static auto registry = c10::RegisterOperators().op(
-    FunctionSchema(
-        "_c10_experimental::Cast",
-        "",
-        (std::vector<c10::Argument>{
-            c10::Argument("input"),
-            c10::Argument("output"),
-            c10::Argument("to_dtype", IntType::get()),
-        }),
-        (std::vector<c10::Argument>{})),
+    "_c10_experimental::Cast",
     c10::kernel<decltype(cast_op_cpu), &cast_op_cpu>(),
     c10::dispatchKey(CPUTensorId()));
 

diff --git a/caffe2/operators/experimental/c10/cpu/concat_cpu.cc b/caffe2/operators/experimental/c10/cpu/concat_cpu.cc
@@ -105,16 +105,7 @@ void concat_op_cpu_impl(
 }
 
 static auto registry = c10::RegisterOperators().op(
-    FunctionSchema(
-        "_c10_experimental::Concat",
-        "",
-        (std::vector<c10::Argument>{
-            c10::Argument("inputs", ListType::ofTensors()),
-            c10::Argument("output"),
-            c10::Argument("split_info"),
-            c10::Argument("add", IntType::get()),
-            c10::Argument("add_axis", IntType::get())}),
-        (std::vector<c10::Argument>{})),
+    "_c10_experimental::Concat",
     c10::kernel<
         decltype(concat_op_cpu_impl<float, CPUContext>),
         &concat_op_cpu_impl<float, CPUContext>>(),

diff --git a/caffe2/operators/experimental/c10/cpu/enforce_finite_cpu.cc b/caffe2/operators/experimental/c10/cpu/enforce_finite_cpu.cc
@@ -25,11 +25,7 @@ void enforce_finite_op_impl_cpu(const at::Tensor& input_) {
 }
 
 static auto registry = c10::RegisterOperators().op(
-    FunctionSchema(
-        "_c10_experimental::EnforceFinite",
-        "",
-        (std::vector<c10::Argument>{c10::Argument("input")}),
-        (std::vector<c10::Argument>{})),
+    "_c10_experimental::EnforceFinite",
     c10::kernel<
         decltype(enforce_finite_op_impl_cpu<float>),
         &enforce_finite_op_impl_cpu<float>>(),

diff --git a/caffe2/operators/experimental/c10/cpu/expand_dims_cpu.cc b/caffe2/operators/experimental/c10/cpu/expand_dims_cpu.cc
@@ -55,13 +55,7 @@ class expand_dims_cpu final : public c10::OperatorKernel {
 };
 
 static auto registry = c10::RegisterOperators().op(
-    FunctionSchema(
-        "_c10_experimental::ExpandDims",
-        "",
-        (std::vector<c10::Argument>{c10::Argument("input"),
-                                    c10::Argument("output"),
-                                    c10::Argument("dims", ListType::ofInts())}),
-        (std::vector<c10::Argument>{})),
+    "_c10_experimental::ExpandDims",
     c10::kernel<expand_dims_cpu<float>>(),
     c10::dispatchKey(CPUTensorId()));
 

diff --git a/caffe2/operators/experimental/c10/cpu/fc_cpu.cc b/caffe2/operators/experimental/c10/cpu/fc_cpu.cc
@@ -129,16 +129,7 @@ class fc_op_cpu final : public c10::OperatorKernel {
 };
 
 static auto registry = c10::RegisterOperators().op(
-    FunctionSchema(
-        "_c10_experimental::FullyConnected",
-        "",
-        (std::vector<c10::Argument>{c10::Argument("X"),
-                                    c10::Argument("W"),
-                                    c10::Argument("b"),
-                                    c10::Argument("output"),
-                                    c10::Argument("axis", IntType::get()),
-                                    c10::Argument("axis_w", IntType::get())}),
-        (std::vector<c10::Argument>{})),
+    "_c10_experimental::FullyConnected",
     c10::kernel<fc_op_cpu<float, CPUContext>>(),
     c10::dispatchKey(CPUTensorId()));
 

diff --git a/caffe2/operators/experimental/c10/cpu/filler_cpu.cc b/caffe2/operators/experimental/c10/cpu/filler_cpu.cc
@@ -145,82 +145,27 @@ void uniform_fill_op_cpu_impl(
 
 static auto registry =
     c10::RegisterOperators()
-        .op(FunctionSchema(
-                "_c10_experimental::ConstantFill",
-                "",
-                (std::vector<c10::Argument>{
-                    c10::Argument("inputs", ListType::ofTensors()),
-                    c10::Argument("output"),
-                    c10::Argument("shape", ListType::ofInts()),
-                    c10::Argument("extra_shape", ListType::ofInts()),
-                    c10::Argument("input_as_shape", BoolType::get()),
-                    c10::Argument("dtype", IntType::get()),
-                    c10::Argument("value", NumberType::get())}),
-                (std::vector<c10::Argument>{})),
+        .op("_c10_experimental::ConstantFill",
             c10::kernel<
                 decltype(constant_fill_op_cpu_impl),
                 &constant_fill_op_cpu_impl>(),
             c10::dispatchKey(CPUTensorId()))
-        .op(FunctionSchema(
-                "_c10_experimental::UniformFill",
-                "",
-                (std::vector<c10::Argument>{
-                    c10::Argument("inputs", ListType::ofTensors()),
-                    c10::Argument("output"),
-                    c10::Argument("shape", ListType::ofInts()),
-                    c10::Argument("extra_shape", ListType::ofInts()),
-                    c10::Argument("input_as_shape", BoolType::get()),
-                    c10::Argument("min", FloatType::get()),
-                    c10::Argument("max", FloatType::get())}),
-                (std::vector<c10::Argument>{})),
+        .op("_c10_experimental::UniformFill",
             c10::kernel<
                 decltype(uniform_fill_op_cpu_impl),
                 &uniform_fill_op_cpu_impl>(),
             c10::dispatchKey(CPUTensorId()))
-        .op(FunctionSchema(
-                "_c10_experimental::GivenTensorFill",
-                "",
-                (std::vector<c10::Argument>{
-                    c10::Argument("inputs", ListType::ofTensors()),
-                    c10::Argument("output"),
-                    c10::Argument("shape", ListType::ofInts()),
-                    c10::Argument("extra_shape", ListType::ofInts()),
-                    c10::Argument("input_as_shape", BoolType::get()),
-                    c10::Argument("values"),
-                }),
-                (std::vector<c10::Argument>{})),
+        .op("_c10_experimental::GivenTensorFill",
             c10::kernel<
                 decltype(given_tensor_fill_op_cpu_impl<float, CPUContext>),
                 &given_tensor_fill_op_cpu_impl<float, CPUContext>>(),
             c10::dispatchKey(CPUTensorId()))
-        .op(FunctionSchema(
-                "_c10_experimental::GivenTensorIntFill",
-                "",
-                (std::vector<c10::Argument>{
-                    c10::Argument("inputs", ListType::ofTensors()),
-                    c10::Argument("output"),
-                    c10::Argument("shape", ListType::ofInts()),
-                    c10::Argument("extra_shape", ListType::ofInts()),
-                    c10::Argument("input_as_shape", BoolType::get()),
-                    c10::Argument("values"),
-                }),
-                (std::vector<c10::Argument>{})),
+        .op("_c10_experimental::GivenTensorIntFill",
             c10::kernel<
                 decltype(given_tensor_fill_op_cpu_impl<int, CPUContext>),
                 &given_tensor_fill_op_cpu_impl<int, CPUContext>>(),
             c10::dispatchKey(CPUTensorId()))
-        .op(FunctionSchema(
-                "_c10_experimental::GivenTensorInt64Fill",
-                "",
-                (std::vector<c10::Argument>{
-                    c10::Argument("inputs", ListType::ofTensors()),
-                    c10::Argument("output"),
-                    c10::Argument("shape", ListType::ofInts()),
-                    c10::Argument("extra_shape", ListType::ofInts()),
-                    c10::Argument("input_as_shape", BoolType::get()),
-                    c10::Argument("values"),
-                }),
-                (std::vector<c10::Argument>{})),
+        .op("_c10_experimental::GivenTensorInt64Fill",
             c10::kernel<
                 decltype(given_tensor_fill_op_cpu_impl<int, CPUContext>),
                 &given_tensor_fill_op_cpu_impl<int, CPUContext>>(),

diff --git a/caffe2/operators/experimental/c10/cpu/flatten_cpu.cc b/caffe2/operators/experimental/c10/cpu/flatten_cpu.cc
@@ -27,13 +27,7 @@ void flatten_op_cpu_impl(
 }
 
 static auto registry = c10::RegisterOperators().op(
-    FunctionSchema(
-        "_c10_experimental::Flatten",
-        "",
-        (std::vector<c10::Argument>{c10::Argument("input"),
-                                    c10::Argument("output"),
-                                    c10::Argument("axis", IntType::get())}),
-        (std::vector<c10::Argument>{})),
+    "_c10_experimental::Flatten",
     c10::kernel<
         decltype(flatten_op_cpu_impl<float, CPUContext>),
         &flatten_op_cpu_impl<float, CPUContext>>(),

diff --git a/caffe2/operators/experimental/c10/cpu/mul_cpu.cc b/caffe2/operators/experimental/c10/cpu/mul_cpu.cc
@@ -72,16 +72,7 @@ void mul_op_cpu_impl(
 }
 
 static auto registry = c10::RegisterOperators().op(
-    FunctionSchema(
-        "_c10_experimental::Mul",
-        "",
-        (std::vector<c10::Argument>{
-            c10::Argument("input1"),
-            c10::Argument("input2"),
-            c10::Argument("output"),
-            c10::Argument("legacy_broadcast", BoolType::get()),
-            c10::Argument("axis", IntType::get())}),
-        (std::vector<c10::Argument>{})),
+    "_c10_experimental::Mul",
     c10::kernel<decltype(mul_op_cpu_impl<float>), &mul_op_cpu_impl<float>>(),
     c10::dispatchKey(CPUTensorId()));
 

diff --git a/caffe2/operators/experimental/c10/cpu/relu_cpu.cc b/caffe2/operators/experimental/c10/cpu/relu_cpu.cc
@@ -41,12 +41,7 @@ void relu_op_cpu_impl(
 }
 
 static auto registry = c10::RegisterOperators().op(
-    FunctionSchema(
-        "_c10_experimental::Relu",
-        "",
-        (std::vector<c10::Argument>{c10::Argument("input"),
-                                    c10::Argument("output")}),
-        (std::vector<c10::Argument>{})),
+    "_c10_experimental::Relu",
     c10::kernel<decltype(relu_op_cpu_impl<float>), &relu_op_cpu_impl<float>>(),
     c10::dispatchKey(CPUTensorId()));
 

diff --git a/caffe2/operators/experimental/c10/cpu/sigmoid_cpu.cc b/caffe2/operators/experimental/c10/cpu/sigmoid_cpu.cc
@@ -24,12 +24,7 @@ void sigmoid_op_cpu_impl(
 }
 
 static auto registry = c10::RegisterOperators().op(
-    FunctionSchema(
-        "_c10_experimental::Sigmoid",
-        "",
-        (std::vector<c10::Argument>{c10::Argument("input"),
-                                    c10::Argument("output")}),
-        (std::vector<c10::Argument>{})),
+    "_c10_experimental::Sigmoid",
     c10::kernel<
         decltype(sigmoid_op_cpu_impl<float>),
         &sigmoid_op_cpu_impl<float>>(),

diff --git a/caffe2/operators/experimental/c10/cpu/sigmoid_cross_entropy_with_logits_cpu.cc b/caffe2/operators/experimental/c10/cpu/sigmoid_cross_entropy_with_logits_cpu.cc
@@ -71,16 +71,7 @@ void sigmoid_cross_entropy_with_logits_op_cpu_impl(
 }
 
 static auto registry = c10::RegisterOperators().op(
-    FunctionSchema(
-        "_c10_experimental::SigmoidCrossEntropyWithLogits",
-        "",
-        (std::vector<c10::Argument>{
-            c10::Argument("input1"),
-            c10::Argument("input2"),
-            c10::Argument("output"),
-            c10::Argument("log_D_trick", BoolType::get()),
-            c10::Argument("unjoined_lr_loss", BoolType::get())}),
-        (std::vector<c10::Argument>{})),
+    "_c10_experimental::SigmoidCrossEntropyWithLogits",
     c10::kernel<
         decltype(sigmoid_cross_entropy_with_logits_op_cpu_impl),
         &sigmoid_cross_entropy_with_logits_op_cpu_impl>(),

diff --git a/caffe2/operators/experimental/c10/cpu/sparse_lengths_sum_cpu.cc b/caffe2/operators/experimental/c10/cpu/sparse_lengths_sum_cpu.cc
@@ -82,14 +82,7 @@ void sparse_lengths_sum_op_cpu(
 }
 
 static auto registry = c10::RegisterOperators().op(
-    FunctionSchema(
-        "_c10_experimental::SparseLengthsSum",
-        "",
-        (std::vector<c10::Argument>{c10::Argument("data"),
-                                    c10::Argument("indices"),
-                                    c10::Argument("lengths"),
-                                    c10::Argument("output")}),
-        (std::vector<c10::Argument>{})),
+    "_c10_experimental::SparseLengthsSum",
     c10::kernel<
         decltype(sparse_lengths_sum_op_cpu),
         &sparse_lengths_sum_op_cpu>(),

diff --git a/caffe2/operators/experimental/c10/cpu/stop_gradient_cpu.cc b/caffe2/operators/experimental/c10/cpu/stop_gradient_cpu.cc
@@ -20,12 +20,7 @@ void stop_gradient_op_cpu_impl(
 }
 
 static auto registry = c10::RegisterOperators().op(
-    FunctionSchema(
-        "_c10_experimental::StopGradient",
-        "",
-        (std::vector<c10::Argument>{c10::Argument("input"),
-                                    c10::Argument("output")}),
-        (std::vector<c10::Argument>{})),
+    "_c10_experimental::StopGradient",
     c10::kernel<
         decltype(stop_gradient_op_cpu_impl<float>),
         &stop_gradient_op_cpu_impl<float>>(),