8000 Return type annotation of `Tensor.long()` etc is not narrowed down to dtype-specific names `LongTensor` etc · Issue #148552 · pytorch/pytorch · GitHub
[go: up one dir, main page]
More Web Proxy on the site http://driver.im/
Skip to content
Return type annotation of Tensor.long() etc is not narrowed down to dtype-specific names LongTensor etc #148552
Open
@lkct

Description

@lkct

🐛 Describe the bug

Sometimes when we want dtype-specific typing for Tensor, e.g. LongTensor, the conversion from Tensor.long() (and similar methods) still returns Tensor, so the following code fails mypy.

from torch import LongTensor, Tensor


def foo(x: Tensor) -> LongTensor:
    return x.long()
a.py:5: error: Incompatible return value type (got "Tensor", expected "LongTensor")  [return-value]

In fact, it seems no function/method is annotated with LongTensor, and the only way to obtain LongTensor for typing is to call the constructor LongTensor() (and, ofc, typing.cast).

A simple solution is to change the dtype conversion methods to return corresponding class names.

A more comprehensive solution could be to make Tensor generic on dtype (similar to numpy), so that methods with dtype= kwarg can also work. However, this might involve too much work.

Yet I would also like to confirm if there's a specific reason that we shouldn't fix it in either way. Still taking LongTensor as an example, it appears many times in the docstrings (e.g. argmax should be another way to create a LongTensor), but never really used in type annotations, so I wonder if it is designed so on purpose.

Versions

nightly torch-2.7.0.dev20250305

cc @ezyang @malfet @xuzhao9 @gramster

Metadata

Metadata

Assignees

No one assigned

    Labels

    module: typingRelated to mypy type annotationstriagedThis 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

      0