How do I add overloads for custom ops defined via torch.Library?`

torch.compile
1

I am trying to use the torch.library API to create my own torch custom operators. The API seems to be missing an option to allow overloaded operators.

For example, suppose I define a Op like this:

torch.library.define("my_lib::my_op", "(Tensor x, float y) -> Tensor")

I want to add an overloaded version that takes two Tensor arguments instead.

torch.library.???("my_lib::my_op", "(Tensor x, Tensor y) -> Tensor", overload_name="tensor_version")

How can I do this?

2

Turned out to be something quite simple. I had to do this, but this wasn’t obvious from torch.library API. The documentation may need to include this.

torch.library.define("my_lib::my_op.<overload_name>", "(Tensor x, float y) -> Tensor")

Now, I need to figure out how to add “out” variants for these ops :slight_smile: