Update on "Improved CustomOp API with schema inference"

zou3519 · zou3519 · commit 0833eda8b1fd · 2023-04-27T14:27:53.000-07:00
This PR changes the CustomOp API. There are now two ways to create a
CustomOp object.

Method 1: with no schema string. We will infer what the schema string is
from your type annotations

```py
custom_op("customlib::foo")
def foo(x: Tensor) -&gt; Tensor:
    ...
```

Method 2: with a schema string, if the inference doesn't work well.

```py
custom_op("customlib::foo", "(Tensor x) -&gt; Tensor")
def foo(x):
    ...
```

Some details:
- We support most combinations of {Tensor, Number, int, float, bool} and
{Optional[typ], Tuple[typ, ...]} as inputs. The combinations we support are mostly
from me reading native_functions.yaml.
- We support only Tensor or Tuple of Tensor of fixed size returns.
- A lot of this PR is input validation for both of the above two
methods. For example, when a user provides a manual schema string, then
their function must not have any type annotations and the number of args
and arg names must match the schema.

Test Plan:
- new tests

[ghstack-poisoned]
diff --git a/test/test_meta.py b/test/test_meta.py
@@ -26,7 +26,7 @@
     OpDTypes,
 )
 from torch.testing._internal.common_methods_invocations import op_db
-from torchgen.yaml import YamlLoader
+from torchgen.yaml_utils import YamlLoader
 from torchgen.model import OperatorName
 
 import sys

Original file line number	Diff line number	Diff line change
`@@ -26,7 +26,7 @@`
`26`	`26`	`OpDTypes,`
`27`	`27`	`)`
`28`	`28`	`from torch.testing._internal.common_methods_invocations import op_db`
`29`		`-from torchgen.yaml import YamlLoader`
	`29`	`+from torchgen.yaml_utils import YamlLoader`
`30`	`30`	`from torchgen.model import OperatorName`
`31`	`31`
`32`	`32`	`import sys`