PyTorch Parameter.__new__ is misidentified as Tensor.__new__

[prashantraina]

Version: VS Code extension v1.13.1 Type checking level: basic

In PyTorch, torch.nn.Parameter is a subclass of torch.Tensor.
For some reason, basedpyright thinks the constructor of Parameter returns a Tensor.
E.g.

# pip install torch
import torch

x: torch.nn.Parameter

x = torch.nn.Parameter(torch.zeros(3))
# basedpyright complains that I'm assigning a torch.Tensor to x

Using Go To Definition on torch.nn.Parameter.__new__ leads to the stub for torch.Tensor.__new__. But using Go To Definition on torch.nn.Parameter correctly leads to the class definition of Parameter, which has an override of __new__ that returns Parameter.

I know ML libraries often have wrong or missing types, but this seems like a language server issue.
Thanks for all you do, btw.

[KotlinIsland]

Is this the stub you are referencing? https://github.com/pytorch/pytorch/blob/main/torch%2Fnn%2Fparameter.pyi

[prashantraina]

Actually no. Small correction, Go to Definition on torch.nn.Parameter.__new__ took me to TensorBase.__new__. After pip installing, the stub file is located in site-packages/torch/_C/__init__.pyi around line 1880. The closest equivalent in their Github repo is https://github.com/pytorch/pytorch/blob/main/torch/_C/__init__.pyi.in The subclass' actual __new__ is defined in https://github.com/pytorch/pytorch/blob/main/torch/nn/parameter.py , line 31, which is installed as site-packages/torch/nn/parameter.py

The main problem I'm facing is that the Parameter constructor is treated as returning a Tensor, which doesn't happen in PyLance.
FYI PyLance also goes to TensorBase.__new__ when Ctrl+clicking on torch.nn.Parameter.__new__.

[KotlinIsland]

Thanks, if a __new__ is added to the stub for Parameter then everything works correctly. we will look into this more