pytorch - 💡(How to fix) Fix register_buffer-created attributes and Parameter.data incorrectly typed as Tensor | Module

🐛 Describe the bug

Module.__getattr__ returns Union[Tensor, Module], which causes false-positive type errors when accessing registered buffers, parameters, and their .data attribute.

Repro 1 — register_buffer produces Tensor | Module

import torch
from torch import nn

class MyModel(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.register_buffer(running_mean, torch.zeros(4))

    def update(self, x: torch.Tensor) -> None:
        self.running_mean.add_(x.mean(dim=0))
        # ty: call-non-callable

Repro 2 — param.data.fill_() reports invalid-argument-type

import torch
from torch import nn

class MyModel(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.weight = nn.Parameter(torch.randn(4, 4))

    def reset(self) -> None:
        self.weight.data.fill_(0)
        # ty: invalid-argument-type

Root cause

torch/nn/modules/module.py:1954 — __getattr__ returns Union[Tensor, Module]: All buffers, parameters, and submodules accessed via self.<name> are typed as Tensor | Module.

Expected behavior

Type checkers should infer correct types for registered buffers and .data on Parameter.

Environment

• PyTorch 2.11, Python 3.14, ty (astral.sh type checker)

cc @albanD @mruberry @jbschlosser @walterddr @mikaylagawarecki @lolpack @maggiemoss @ndmitchell @kinto0 @samwgoldman