Skip to content

[custom_op] Better error message on unsupported Tuple types. #129370

Description

@zou3519
import torch
from torch import Tensor
from typing import Tuple

@torch.library.custom_op("mylib::foo", mutates_args={})
def foo(x: Tensor, y: Tuple[int, int]) -> Tensor:
    return x * y[0] * y[1]

gives:

ValueError: infer_schema(func): Parameter y has unsupported type typing.Tuple[int, int]. The valid typ
es are: dict_keys([<class 'torch.Tensor'>, typing.Optional[torch.Tensor], typing.Sequence[torch.Tensor
], typing.List[torch.Tensor], typing.Sequence[typing.Optional[torch.Tensor]], typing.List[typing.Optio
nal[torch.Tensor]], <class 'int'>, typing.Optional[int], typing.Sequence[int], typing.List[int], typin
g.Optional[typing.Sequence[int]], typing.Optional[typing.List[int]], <class 'float'>, typing.Optional[
float], typing.Sequence[float], typing.List[float], typing.Optional[typing.Sequence[float]], typing.Op
tional[typing.List[float]], <class 'bool'>, typing.Optional[bool], typing.Sequence[bool], typing.List[
bool], typing.Optional[typing.Sequence[bool]], typing.Optional[typing.List[bool]], <class 'str'>, typi
ng.Optional[str], typing.Union[int, float, bool], typing.Union[int, float, bool, NoneType], typing.Seq
uence[typing.Union[int, float, bool]], typing.List[typing.Union[int, float, bool]], <class 'torch.dtyp
e'>, typing.Optional[torch.dtype], <class 'torch.device'>, typing.Optional[torch.device]]). Got func w
ith signature (x: torch.Tensor, y: Tuple[int, int]) -> torch.Tensor)

What we really meant to say was: "We don't support Tuple arguments. Please try to use a List[int] instead".

Bonus points if the error message can recommend the correct type. That is, I have seen users wonder why a List[Tensor, Tensor, Optional[Tensor]] doesn't work: our error message should recommend List[Optional[Tensor]] in this case.

cc @ezyang @anijain2305 @chauhang @bdhirsh

Metadata

Metadata

Assignees

Labels

actionablemodule: custom-operatorscustom operators, custom ops, custom-operators, custom-opsmodule: pt2-dispatcherPT2 dispatcher-related issues (e.g., aotdispatch, functionalization, faketensor, custom-op,oncall: pt2triagedThis issue has been looked at a team member, and triaged and prioritized into an appropriate module

Type

No type

Projects

No projects

Milestone

No milestone

Relationships

None yet

Development

No branches or pull requests

Issue actions