Use compact shape comments throughout the numerical corpus

Summary:
Replace the long `type =` and `.args[N]` spellings in the NumPy and SciPy demos with `= payload` and `.N = payload`. Keep all visible type annotations and shape contracts unchanged while making array carriers and selectors easier to read.

Update the corpus's notation notes and retain the earlier spellings as supported parser aliases.

Differential Revision: D127615303

Author: stroxler

tensor-shapes/comment-shape-corpus/README.md

@@ -22,6 +22,14 @@ Start with [the SciPy demo](scipy/README.md):
 - [Shape manipulation](numpy/shape_base.pyi): repetition tuples, symbolic
   multiplication, and an existing DSL-backed return transformation.
 
+## Compact notation
+
+`# shaped: payload` uses compatibility semantics; `# shaped: = payload` replaces
+the attached type component. `.0 = payload` selects its first written type
+argument, and paths such as `.0.1` reach nested arguments. The older `type =`
+and `.args[N]` forms remain accepted aliases, but the examples use the shorter
+spellings. Placement, shape contracts, and supported API cases are unchanged.
+
 ## Run
 
 Use a Pyrefly build containing shape-comment support. Only NumPy's canonical

tensor-shapes/comment-shape-corpus/numpy/README.md

@@ -54,7 +54,8 @@ not every rejection indicates that NumPy itself would raise an exception.
 - `roll` demonstrates one whole-shape variable and unchanged dtype typing.
 - `diag` uses independent shape declarations on each overload.
 - Matrix contractions connect the dimensions of separate arguments.
-- Comments directly on tuple carriers keep dtype arguments readable.
+- Comments directly on tuple carriers use `# shaped: = payload` and keep dtype
+  arguments readable. A whole-array comment can use `.0 = payload` instead.
 - `tile` turns an ordinary `tuple[int, int]` argument into an `IntTuple` and
   multiplies symbolic dimensions without changing the visible parameter type.
 - `expand_dims` demonstrates a DSL call inside the return carrier comment. Its

tensor-shapes/comment-shape-corpus/numpy/numeric.pyi

@@ -11,45 +11,45 @@ _Scalar = TypeVar("_Scalar", bound=np.generic)
 # shapes: M, N
 def outer(
     a: np.ndarray[
-        tuple[int],  # shaped: type = [M]
+        tuple[int],  # shaped: = [M]
         np.dtype[np.float64],
     ],
     b: np.ndarray[
-        tuple[int],  # shaped: type = [N]
+        tuple[int],  # shaped: = [N]
         np.dtype[np.float64],
     ],
     out: (
         np.ndarray[
-            tuple[int, int],  # shaped: type = [M, N]
+            tuple[int, int],  # shaped: = [M, N]
             np.dtype[np.float64],
         ]
         | None
     ) = None,
 ) -> np.ndarray[
-    tuple[int, int],  # shaped: type = [M, N]
+    tuple[int, int],  # shaped: = [M, N]
     np.dtype[np.float64],
 ]: ...
 
 # shapes: M, N, P
 def tensordot(
     a: np.ndarray[
-        tuple[int, int],  # shaped: type = [M, N]
+        tuple[int, int],  # shaped: = [M, N]
         np.dtype[np.float64],
     ],
     b: np.ndarray[
-        tuple[int, int],  # shaped: type = [N, P]
+        tuple[int, int],  # shaped: = [N, P]
         np.dtype[np.float64],
     ],
     axes: Literal[1],
 ) -> np.ndarray[
-    tuple[int, int],  # shaped: type = [M, P]
+    tuple[int, int],  # shaped: = [M, P]
     np.dtype[np.float64],
 ]: ...
 
 # shapes: *Shape
 def roll(
-    a: np.ndarray[tuple[int, ...], np.dtype[_Scalar]],  # shaped: .args[0] = Shape
+    a: np.ndarray[tuple[int, ...], np.dtype[_Scalar]],  # shaped: .0 = Shape
     shift: int | tuple[int, ...],
     axis: int | tuple[int, ...] | None = None,
-) -> np.ndarray[tuple[int, ...], np.dtype[_Scalar]]:  # shaped: .args[0] = Shape
+) -> np.ndarray[tuple[int, ...], np.dtype[_Scalar]]:  # shaped: .0 = Shape
     ...

tensor-shapes/comment-shape-corpus/numpy/shape_base.pyi

@@ -12,38 +12,38 @@ _Scalar = TypeVar("_Scalar", bound=np.generic)
 # shapes: *Shape
 def expand_dims(
     a: np.ndarray[
-        tuple[int, ...],  # shaped: type = Shape
+        tuple[int, ...],  # shaped: = Shape
         np.dtype[_Scalar],
     ],
     axis: Literal[0],
 ) -> np.ndarray[
-    tuple[int, ...],  # shaped: type = expand_dims_shape(Shape, 0)
+    tuple[int, ...],  # shaped: = expand_dims_shape(Shape, 0)
     np.dtype[_Scalar],
 ]: ...
 
 # shapes: M, N, R, C
 def tile(
     A: np.ndarray[
-        tuple[int, int],  # shaped: type = [M, N]
+        tuple[int, int],  # shaped: = [M, N]
         np.dtype[_Scalar],
     ],
-    reps: tuple[int, int],  # shaped: type = [R, C]
+    reps: tuple[int, int],  # shaped: = [R, C]
 ) -> np.ndarray[
-    tuple[int, int],  # shaped: type = [M * R, N * C]
+    tuple[int, int],  # shaped: = [M * R, N * C]
     np.dtype[_Scalar],
 ]: ...
 
 # shapes: M, N, P, Q
 def kron(
     a: np.ndarray[
-        tuple[int, int],  # shaped: type = [M, N]
+        tuple[int, int],  # shaped: = [M, N]
         np.dtype[np.float64],
     ],
     b: np.ndarray[
-        tuple[int, int],  # shaped: type = [P, Q]
+        tuple[int, int],  # shaped: = [P, Q]
         np.dtype[np.float64],
     ],
 ) -> np.ndarray[
-    tuple[int, int],  # shaped: type = [M * P, N * Q]
+    tuple[int, int],  # shaped: = [M * P, N * Q]
     np.dtype[np.float64],
 ]: ...

tensor-shapes/comment-shape-corpus/numpy/twodim_base.pyi

@@ -12,48 +12,48 @@ _Scalar = TypeVar("_Scalar", bound=np.generic)
 @overload
 def diag(
     v: np.ndarray[
-        tuple[int],  # shaped: type = [N]
+        tuple[int],  # shaped: = [N]
         np.dtype[_Scalar],
     ],
     k: Literal[0] = 0,
 ) -> np.ndarray[
-    tuple[int, int],  # shaped: type = [N, N]
+    tuple[int, int],  # shaped: = [N, N]
     np.dtype[_Scalar],
 ]: ...
 
 # shapes: N
 @overload
 def diag(
     v: np.ndarray[
-        tuple[int, int],  # shaped: type = [N, N]
+        tuple[int, int],  # shaped: = [N, N]
         np.dtype[_Scalar],
     ],
     k: Literal[0] = 0,
 ) -> np.ndarray[
-    tuple[int],  # shaped: type = [N]
+    tuple[int],  # shaped: = [N]
     np.dtype[_Scalar],
 ]: ...
 
 # shapes: *Batch, M, N
 def tril(
     v: np.ndarray[
-        tuple[int, ...],  # shaped: type = [*Batch, M, N]
+        tuple[int, ...],  # shaped: = [*Batch, M, N]
         np.dtype[_Scalar],
     ],
     k: int = 0,
 ) -> np.ndarray[
-    tuple[int, ...],  # shaped: type = [*Batch, M, N]
+    tuple[int, ...],  # shaped: = [*Batch, M, N]
     np.dtype[_Scalar],
 ]: ...
 
 # shapes: *Batch, M, N
 def triu(
     v: np.ndarray[
-        tuple[int, ...],  # shaped: type = [*Batch, M, N]
+        tuple[int, ...],  # shaped: = [*Batch, M, N]
         np.dtype[_Scalar],
     ],
     k: int = 0,
 ) -> np.ndarray[
-    tuple[int, ...],  # shaped: type = [*Batch, M, N]
+    tuple[int, ...],  # shaped: = [*Batch, M, N]
     np.dtype[_Scalar],
 ]: ...

tensor-shapes/comment-shape-corpus/scipy/README.md

@@ -49,7 +49,7 @@ properties such as positive definiteness, finite values, or valid FFT axes.
 
 ## What to assess
 
-Compare the small `int  # shaped: N` case with array `.args[0]` selectors,
+Compare the small `int  # shaped: N` case with array `.0` selectors,
 comments directly beside tuple shape carriers, and the whole-type replacement
 needed when an `NDArray` alias hides its shape argument. Try changing ordinary
 annotations or adding API cases: the goal is to expose syntax friction, not to
@@ -58,9 +58,9 @@ maximize SciPy coverage.
 ## Syntax observations
 
 - The scalar `n: int  # shaped: N` form is compact.
-- `.args[0]` preserves dtype annotations. Splitting the type lets a shape comment
+- `.0` preserves dtype annotations. Splitting the type lets a shape comment
   sit directly beside `tuple[int, ...]` instead of using a long selector.
-- An `NDArray` alias hides the shape slot, requiring a full `type =` replacement.
+- An `NDArray` alias hides the shape slot, requiring a full `= payload` replacement.
 - `cho_factor` refines both the factor array and its array-valued lower flag.
   Per-component comments express both shapes without repeating the return type.
 - Optional `out` arrays keep their shape comment on the array's carrier line,

tensor-shapes/comment-shape-corpus/scipy/cholesky.pyi

@@ -12,33 +12,33 @@ type _LowerFlag = bool | np.ndarray[tuple[()], np.dtype[np.bool]]
 # shapes: *Batch, N
 def cholesky(
     a: np.ndarray[
-        tuple[int, ...],  # shaped: type = [*Batch, N, N]
+        tuple[int, ...],  # shaped: = [*Batch, N, N]
         np.dtype[np.float64],
     ],
     lower: bool = False,
     overwrite_a: bool = False,
     check_finite: bool = True,
 ) -> np.ndarray[
-    tuple[int, ...],  # shaped: type = [*Batch, N, N]
+    tuple[int, ...],  # shaped: = [*Batch, N, N]
     np.dtype[np.float64],
 ]: ...
 
 # shapes: *Batch, N
 def cho_factor(
     a: np.ndarray[
-        tuple[int, ...],  # shaped: type = [*Batch, N, N]
+        tuple[int, ...],  # shaped: = [*Batch, N, N]
         np.dtype[np.float64],
     ],
     lower: bool = False,
     overwrite_a: bool = False,
     check_finite: bool = True,
 ) -> tuple[
     np.ndarray[
-        tuple[int, ...],  # shaped: type = [*Batch, N, N]
+        tuple[int, ...],  # shaped: = [*Batch, N, N]
         np.dtype[np.float64],
     ],
     np.ndarray[
-        tuple[int, ...],  # shaped: type = Batch
+        tuple[int, ...],  # shaped: = Batch
         np.dtype[np.bool],
     ],
 ]: ...
@@ -47,24 +47,24 @@ def cho_factor(
 @overload
 def cho_solve(
     c_and_lower: tuple[
-        np.ndarray[tuple[int, int], np.dtype[np.float64]],  # shaped: .args[0] = [N, N]
+        np.ndarray[tuple[int, int], np.dtype[np.float64]],  # shaped: .0 = [N, N]
         _LowerFlag,
     ],
-    b: np.ndarray[tuple[int], np.dtype[np.float64]],  # shaped: .args[0] = [N]
+    b: np.ndarray[tuple[int], np.dtype[np.float64]],  # shaped: .0 = [N]
     overwrite_b: bool = False,
     check_finite: bool = True,
-) -> np.ndarray[tuple[int], np.dtype[np.float64]]:  # shaped: .args[0] = [N]
+) -> np.ndarray[tuple[int], np.dtype[np.float64]]:  # shaped: .0 = [N]
     ...
 
 # shapes: N, K
 @overload
 def cho_solve(
     c_and_lower: tuple[
-        np.ndarray[tuple[int, int], np.dtype[np.float64]],  # shaped: .args[0] = [N, N]
+        np.ndarray[tuple[int, int], np.dtype[np.float64]],  # shaped: .0 = [N, N]
         _LowerFlag,
     ],
-    b: np.ndarray[tuple[int, int], np.dtype[np.float64]],  # shaped: .args[0] = [N, K]
+    b: np.ndarray[tuple[int, int], np.dtype[np.float64]],  # shaped: .0 = [N, K]
     overwrite_b: bool = False,
     check_finite: bool = True,
-) -> np.ndarray[tuple[int, int], np.dtype[np.float64]]:  # shaped: .args[0] = [N, K]
+) -> np.ndarray[tuple[int, int], np.dtype[np.float64]]:  # shaped: .0 = [N, K]
     ...

tensor-shapes/comment-shape-corpus/scipy/distance.pyi

@@ -11,42 +11,40 @@ import numpy.typing as npt
 # shapes: M, P, N
 def cdist(
     XA: np.ndarray[
-        tuple[int, int],  # shaped: type = [M, N]
+        tuple[int, int],  # shaped: = [M, N]
         np.dtype[np.float64],
     ],
     XB: np.ndarray[
-        tuple[int, int],  # shaped: type = [P, N]
+        tuple[int, int],  # shaped: = [P, N]
         np.dtype[np.float64],
     ],
     metric: Literal["euclidean"] = "euclidean",
     *,
     out: (
         np.ndarray[
-            tuple[int, int],  # shaped: type = [M, P]
+            tuple[int, int],  # shaped: = [M, P]
             np.dtype[np.float64],
         ]
         | None
     ) = None,
 ) -> np.ndarray[
-    tuple[int, int],  # shaped: type = [M, P]
+    tuple[int, int],  # shaped: = [M, P]
     np.dtype[np.float64],
 ]: ...
 
 # shapes: M, N
 def pdist(
-    X: npt.NDArray[
-        np.float64
-    ],  # shaped: type = np.ndarray[[M, N], np.dtype[np.float64]]
+    X: npt.NDArray[np.float64],  # shaped: = np.ndarray[[M, N], np.dtype[np.float64]]
     metric: Literal["euclidean"] = "euclidean",
     *,
     out: (
         np.ndarray[
-            tuple[int],  # shaped: type = [M * (M - 1) // 2]
+            tuple[int],  # shaped: = [M * (M - 1) // 2]
             np.dtype[np.float64],
         ]
         | None
     ) = None,
 ) -> np.ndarray[
-    tuple[int],  # shaped: type = [M * (M - 1) // 2]
+    tuple[int],  # shaped: = [M * (M - 1) // 2]
     np.dtype[np.float64],
 ]: ...

tensor-shapes/comment-shape-corpus/scipy/fft_helper.pyi

@@ -16,7 +16,7 @@ def fftfreq(
     *,
     xp: None = None,
     device: Literal["cpu"] | None = None,
-) -> np.ndarray[tuple[int], np.dtype[np.float64]]:  # shaped: .args[0] = [N]
+) -> np.ndarray[tuple[int], np.dtype[np.float64]]:  # shaped: .0 = [N]
     ...
 
 # shapes: N
@@ -26,19 +26,19 @@ def rfftfreq(
     *,
     xp: None = None,
     device: Literal["cpu"] | None = None,
-) -> np.ndarray[tuple[int], np.dtype[np.float64]]:  # shaped: .args[0] = [N // 2 + 1]
+) -> np.ndarray[tuple[int], np.dtype[np.float64]]:  # shaped: .0 = [N // 2 + 1]
     ...
 
 # shapes: *Shape
 def fftshift(
-    x: np.ndarray[tuple[int, ...], np.dtype[_Scalar]],  # shaped: .args[0] = Shape
+    x: np.ndarray[tuple[int, ...], np.dtype[_Scalar]],  # shaped: .0 = Shape
     axes: int | tuple[int, ...] | None = None,
-) -> np.ndarray[tuple[int, ...], np.dtype[_Scalar]]:  # shaped: .args[0] = Shape
+) -> np.ndarray[tuple[int, ...], np.dtype[_Scalar]]:  # shaped: .0 = Shape
     ...
 
 # shapes: *Shape
 def ifftshift(
-    x: np.ndarray[tuple[int, ...], np.dtype[_Scalar]],  # shaped: .args[0] = Shape
+    x: np.ndarray[tuple[int, ...], np.dtype[_Scalar]],  # shaped: .0 = Shape
     axes: int | tuple[int, ...] | None = None,
-) -> np.ndarray[tuple[int, ...], np.dtype[_Scalar]]:  # shaped: .args[0] = Shape
+) -> np.ndarray[tuple[int, ...], np.dtype[_Scalar]]:  # shaped: .0 = Shape
     ...