False positive for `torch.register_*`

[AdrianSosic]

Describe the Bug

Hi, I'm not sure if this is a false positive or if this is even supposed to be handled by static type checking, but I'm wondering how to handle the following situation:

import torch
from torch import nn


class LeModule(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.register_buffer("le_tensor", torch.tensor([0, 1, 2]))


m = LeModule()
if m.le_tensor is not None:
    result = m.le_tensor.tolist()  # pyrefly: Expected a callable, got Tensor

Sandbox Link

No response

(Only applicable for extension issues) IDE Information

No response

[yangdanny97]

The type system can't model a function whose side effect is dynamically adding a field to an object.

Might be something we could special-case for Pytorch specifically if this is a very common pattern tho

[AdrianSosic]

Hi @yangdanny97 👋🏼 I already expected this, but thanks for confirming 👌🏼

Yes, I think some form of special handling would be extremely useful since ’register_parameter’ and ’register_buffer’ are pytorch‘s bread and butter. Also, I guess the alternative of explicit attribute creation using a ’nn.Parameter’ would be type-safe but is much less commonly used.