Pyright is stubborn about using union types in Pytorch

[InfiniteSwerve]

In line 15 of the following working code, Pyright produces a type error:

import torch.nn as nn
import torch.nn.functional as F
from typing import List

class iter_error(nn.Module):
  def __init__(self, layer_sizes : List[int]):
      super().__init__()
      self.layers : nn.ModuleList = nn.ModuleList() # explicit annotation 
      self.layer_sizes = layer_sizes
    
      for h_in, h_out in zip(self.layer_sizes[:-1], self.layer_sizes[1:]):
         self.layers.append(nn.Linear(h_in, h_out))
         
  def forward(self, x):
    for i, layer in enumerate(self.layers[:-1]): # type error is signaled inside this enumerate
        x = F.relu(layer(x))
    x = self.layers[-1](x)
    return x

The error:

Argument of type "Module | ModuleList" cannot be assigned to parameter "iterable" of type "Iterable[_T@enumerate]" in function "__init__"
    Type "Module | ModuleList" cannot be assigned to type "Iterable[_T@enumerate]"
      "Module" is incompatible with protocol "Iterable[_T@enumerate]"
        "__iter__" is not present (reportGeneralTypeIssues)     

It looks like Pyright considers self.layers a union type of Module and ModuleList, even if I explicitly annotate otherwise. ModuleList contains both __getitem__ and __iter__ methods, while Module contains neither, so I think this must be the problem.
How can I communicate my intention to Pyright more clearly?

[erictraut]

The expression self.layers is evaluated as type ModuleList, but the expression self.layers[:-1] has the type ModuleList | Module. This comes directly from the ModuleList.__getitem__ method type declaration:

class ModuleList(Module):
    ...
    def __getitem__(self, idx: Union[int, slice]) -> Union[Module, 'ModuleList']: ...

The problem appears to be that __getitem__ is missing an @overload declaration. If you add the following overloads to the pytorch sources, it will work as expected.

    @overload
    def __getitem__(self, idx: int) -> Module: ...
    @overload
    def __getitem__(self, idx: slice) -> ModuleList: ...

You may want to file a bug report or submit a PR in the pytorch project.