Improved CustomOp API with schema inference (#100127)
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) -> Tensor:
...
```
Method 2: with a schema string, if the inference doesn't work well.
```py
@custom_op("customlib::foo", "(Tensor x) -> 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
Pull Request resolved: https://github.com/pytorch/pytorch/pull/100127
Approved by: https://github.com/ezyang
