[custom_op] torch.library.define should be able to auto-infer schema. #129617
Description
we should be able to pass torch.library.define and torch.library.Library.define a function that the schema can be inferred from, like so:
# instead of the following:
m = torch.library.Library("mylib", "FRAGMENT")
m.define("foo_op(Tensor x) -> Tensor")
def foo_impl(x: torch.Tensor) -> torch.Tensor:
return x.sin()
m.impl("foo", "CPU", foo_impl)
# a user should be able to write:
m.define(torch.library.Schema("foo_op", signature=foo_impl, mutates_args={}))
def foo_impl(x: torch.Tensor) -> torch.Tensor:
return x.sin()
m.impl("foo", "CPU", foo_impl")
Motivation:
- torch.library.custom_op doesn't work in all cases. It's not supposed to work in all cases. There's friction when switching between torch.library.custom_op and torch.library.define