[tensor-shapes] A shape-changing op inside a loop gives a loop-carried union that omits a reachable shape, plus a confusing bad-argument-type #5164

Summary

When a loop-carried tensor goes through a shape-changing layer (a stride-2 conv), the type Pyrefly
infers for it is [4, 4] | [1, 1]. In a loop that runs twice (variant D), the tensor really is
[2, 2] at the head of the second iteration, so the inferred union leaves out a shape that occurs
at runtime. Pyrefly then reports errors on code that is correct at runtime. Each of those reports
describes the conv's input parameter as Tensor[[B, 16, 1, 1]], even though Conv2d.forward
accepts any H and W.

Shape-preserving layers in a loop work fine (variant C).

Environment

• Python 3.12.8, macOS; pyrefly==1.3.2, pyrefly-torch-stubs==1.3.2, pyrefly-shape-extensions==1.3.2, all installed in one venv
• Command: pyrefly check repro.py --preset default --python-interpreter-path .venv/bin/python
• Run from a directory with no Pyrefly configuration and no configuration Pyrefly can migrate from
  another checker (no pyrefly.toml, no [tool.pyrefly] in pyproject.toml, no mypy or pyright config)
• Note: running the venv's pyrefly executable without activating the venv or passing
  --python-interpreter-path picked up a different Python on our machine and failed to resolve
  shape_extensions. Activate the venv or pass the interpreter path to reproduce these results.

Repro

from __future__ import annotations
from torch import Tensor, nn
from shape_extensions import IntVar
from typing import reveal_type

class Down(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.down = nn.Conv2d(16, 16, 3, stride=2, padding=1)   # halves H and W

    # D: nonempty loop. At runtime y is [4, 4] at the first loop head, [2, 2] at the second,
    #    and [1, 1] after the loop. Correct at runtime.
    def loop_two[B: IntVar](self, x: Tensor[[B, 16, 8, 8]]) -> Tensor[[B, 16, 1, 1]]:
        y = self.down(x)
        for _ in range(2):
            reveal_type(y)
            y = self.down(y)
        return y

    # E: loop count known only at runtime (n=1 gives [2, 2] after the loop).
    def loop_n[B: IntVar](self, x: Tensor[[B, 16, 8, 8]], n: int) -> Tensor[[B, 16, 4, 4]]:
        y = self.down(x)
        for _ in range(n):
            y = self.down(y)
        reveal_type(y)
        return y

    # A: loop that never runs. Shows the confusing diagnostic on code that is correct at runtime.
    def loop_zero[B: IntVar](self, x: Tensor[[B, 16, 8, 8]]) -> Tensor[[B, 16, 4, 4]]:
        y = self.down(x)
        for _ in range(0):
            y = self.down(y)
        return y

class Same(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.conv = nn.Conv2d(16, 16, 3, padding=1)   # shape-preserving

    # C: shape-preserving op inside a loop: works.
    def loop_same[B: IntVar](self, x: Tensor[[B, 16, 8, 8]]) -> Tensor[[B, 16, 8, 8]]:
        y = x
        for _ in range(3):
            y = self.conv(y)
        reveal_type(y)
        return y

At runtime (torch 2.10), variant D's y has shape (1, 16, 4, 4) at the first loop head,
(1, 16, 2, 2) at the second, and (1, 16, 1, 1) after the loop.

Actual output (edited excerpt)

 INFO revealed type: Tensor[IntTuple[B, 16, 1, 1]] | Tensor[IntTuple[B, 16, 4, 4]]           # D, loop head (line 16)
ERROR Argument `Tensor[IntTuple[B, 16, 1, 1]] | Tensor[IntTuple[B, 16, 4, 4]]` is not assignable to parameter `input` with type `Tensor[IntTuple[B, 16, 1, 1]]` in function `torch.nn.Conv2d.forward` [bad-argument-type]   # D (line 17)
ERROR Returned type `Tensor[IntTuple[B, 16, 1, 1]] | Tensor[IntTuple[B, 16, 4, 4]]` is not assignable to declared return type `Tensor[IntTuple[B, 16, 1, 1]]` [bad-return]   # D (line 18)
ERROR Argument `... | ...` is not assignable to parameter `input` with type `Tensor[IntTuple[B, 16, 1, 1]]` ... [bad-argument-type]   # E (line 24)
 INFO revealed type: Tensor[IntTuple[B, 16, 1, 1]] | Tensor[IntTuple[B, 16, 4, 4]]           # E, after loop (line 25)
ERROR Returned type `... | ...` is not assignable to declared return type `Tensor[IntTuple[B, 16, 4, 4]]` [bad-return]   # E (line 26)
ERROR Argument `... | ...` is not assignable to parameter `input` with type `Tensor[IntTuple[B, 16, 1, 1]]` ... [bad-argument-type]   # A (line 32)
ERROR Returned type `... | ...` is not assignable to declared return type `Tensor[IntTuple[B, 16, 4, 4]]` [bad-return]   # A (line 33)
 INFO revealed type: Tensor[IntTuple[B, 16, 8, 8]]                                           # C (line 45)
 INFO 6 errors

Observations

1. The loop-head type omits a reachable shape (variant D). In for _ in range(2), y is
   [B, 16, 2, 2] at the head of the second iteration, but the revealed type at the loop head is
   [1, 1] | [4, 4]. Variant E (runtime count, where n = 1 leaves y at [2, 2]) gets the same
   union. So the union isn't the precise set of shapes y takes in this loop.
2. The bad-argument-type inside the loop is confusing (variants A, D, E). Conv2d.forward
   takes Tensor[[B, 16, @_, @_]], but the message reports the parameter type as
   Tensor[[B, 16, 1, 1]], apparently a solved instance from one union member. Since any H/W is
   accepted, it's unclear from the message why the call is rejected. Variant A (range(0), body
   never runs) shows this message on code where the loop body is never executed at all.
3. Shape-preserving loops (variant C) infer the exact shape, so this is specific to loop-carried
   values whose shape changes.

Suggested improvement

• For loop-carried tensors whose dimensions change across iterations, widen the changing
  dimensions to unknown (@_ / Int[int]), or include every shape reached before the fixed point,
  instead of a union that skips intermediate shapes.
• When reporting bad-argument-type against a generic forward, show the declared parameter type
  (Tensor[[B, 16, @_, @_]]) or say which union member failed, rather than a solved instance.

Note on 1.4.0.dev3

Tested on 1.3.2 (latest stable). On pyrefly==1.4.0.dev3 with pyrefly-torch-stubs==1.4.0.dev3, the shape checks in this repro don't fire at all, so it can't show the behavior either way. For example, nn.Conv2d(16, 16, 3) applied to a Tensor[[B, 16, 8, 8]] input returns Tensor[[@_, 16, @_ - 2, @_ - 2]] on 1.4.0.dev3, where 1.3.2 gives Tensor[[B, 16, 6, 6]]. The batch and spatial dimensions are unknown after one conv, so no shape errors are reported.