Unable to return Tensor or List[Tensor] as Any type from TorchScript function

[yf225]

🐛 Bug

I tried this:

from typing import Any, List
import torch

class TestModule(torch.nn.Module):
  def __init__(self):
    super().__init__()

  def forward(self, input: torch.Tensor) -> Any:
    ret: List[torch.Tensor] = []
    if input.shape[0] == 1:
      return input
    else:
      ret.append(input)
      return ret

m = TestModule()
m_scripted = torch.jit.script(m)

and it throws:

Traceback (most recent call last):
  File "test_yf225.py", line 17, in <module>
    m_scripted = torch.jit.script(m)
  File "/data/miniconda3/envs/master_nightly/lib/python3.7/site-packages/torch/jit/__init__.py", line 1516, in script
    return torch.jit._recursive.create_script_module(obj, torch.jit._recursive.infer_methods_to_compile)
  File "/data/miniconda3/envs/master_nightly/lib/python3.7/site-packages/torch/jit/_recursive.py", line 318, in create_script_module
    return create_script_module_impl(nn_module, concrete_type, stubs_fn)
  File "/data/miniconda3/envs/master_nightly/lib/python3.7/site-packages/torch/jit/_recursive.py", line 376, in create_script_module_impl
    create_methods_from_stubs(concrete_type, stubs)
  File "/data/miniconda3/envs/master_nightly/lib/python3.7/site-packages/torch/jit/_recursive.py", line 292, in create_methods_from_stubs
    concrete_type._create_methods(defs, rcbs, defaults)
RuntimeError: it != all_nodes.end() INTERNAL ASSERT FAILED at "/pytorch/torch/csrc/jit/ir/ir.cpp":1806, please report a bug to PyTorch.

Expected behavior

From my understanding, returning Any type from TorchScript function should work.

Environment

PyTorch nightly build

Additional context

One of the DPER3 modules (https://fburl.com/diffusion/bkbuoid2) used by the PyPer Ads ctr_mbl_feed model can return either torch.Tensor or List[torch.Tensor] based on a constructor argument. Without Any support here, we will need to split the module into two modules (corresponding to two different return types), which will cause a lot of confusion and downstream refactoring, slowing down the progress on scripting PyPer Ads models.

cc. @SplitInfinity @wanchaol @suo

cc @suo @gmagogsfm

[gmagogsfm]

I think there are two issues here:

1. two branches of if actually return two different data types. input is torch.Tensor while ret is List[torch.Tensor]. This is not allowed. The type Any refers to a single unspecified type rather than a varying type.

2. Error message isn't pretty and not meaningful

[SplitInfinity]

I'm a bit confused.

two branches of if actually return two different data types. input is torch.Tensor while ret is List[torch.Tensor]. This is not allowed.

In the TorchScript language reference, the only statement that resembles this is one that says that unlike Python, each variable in TorchScript function must have a single static type. So I was under the impression that assigning values of different types to one variable in different branches of an if statement is not allowed, but returning values of different types from each branch is allowed. This usually doesn't work because of how JIT performs type checking for return types.

If the function has no return type annotation, the compiler attempts to unify/merge the different return types and pick the most specific type that supertypes them all. If the function has a return type annotation, these values must be subtypes of that annotated type and the code won't compile if not. Any is treated like any other return type annotation; instead of attempting to find the most specific type that supertypes all return types, the compiler checks that all return types are subtypes of the annotated return type (Any). This check, of course, never fails.

In #42259, I modified the checking of return types in the IR generator so that it would try to find the most specific supertype of all return types even when the annotation is Any. If the meaning of Any as a return type is truly a single unspecified type rather than a varying type, then trying to find the aforementioned most specific supertype is the right thing to do. But in the discussion on the PR, someone commented that doing this breaks backward compatibility (indicating that there is working code out there that can break if we try to unify all return types in a function with an Any return type). During our meeting this week, someone also commented that if the return type is annotated as Any, then all we have to do is check that every return type converts to Any, and we don't need to unify/merge any types. All of this gives me the impression that Any truly means a varying type.

And finally, if the meaning of Any in TorchScript is a single, unspecified type, what is its usefulness? In Python, it means literally any type, so most legal Python code that uses Any as a return type probably won't be legal TorchScript. I know TorchScript is a subset of Python but the point I'm trying to make is that this makes it harder for the same codebase to be run in Python and in JIT. If Any did indeed refer to a single, unspecified type, the author of that code would have used the annotation for that type instead. If this is truly what Any means as a return annotation, I think we should replace it with TypeVar.

[gmagogsfm]

Thanks @SplitInfinity for the clarification.

In last meeting, I asked specifically whether Any refers to a single unspecified type or a varying type in the same meeting and someone mentioned that it refers to a single type. Now that I think more about it, the only use case of it is to represent an unknown or uninferrable type, which may or may not be specified later. However, you made a very good point that it doesn't match Python's definition of Any.

Just scanned through code/comment/document, didn't find a good answer. @suo any insight?

[yf225]

Thanks a lot everyone and yes would love to know if I should expect the feature to be possible / available in the near future, so that I can also communicate to our team as well. cc. @suo

[wanchaol]

cc @suo for GraphModule engagement.

[SplitInfinity]

Does #42259 address this?

[yf225]

Does #42259 address this?

Yes it does, thanks a lot for the awesome feature @SplitInfinity !