[tensor-shapes] jax.numpy.reshape: support -1 dimensions

tensor-shapes/jax-pyrefly-stubs/jax-stubs/_array.pyi

@@ -68,10 +68,6 @@ type _IntegerSequence = (
     | tuple[SupportsIndex | _IntegerSequence, ...]
 )
 type _BasicIndex = SupportsIndex | slice | _IntegerSequence | None | EllipsisType
-# The trailing `None` is not a legal argument to `reshape`. It is present because
-# an `int | tuple[int, ...]` parameter cannot be iterated inside a DSL function
-# after narrowing with `is_int_value` alone. See `reshape_shape`, which rejects it.
-type _NewShape = int | tuple[int, ...] | None
 
 class Array[Shape: _Shape = _Shape]:
     shape: Shape
@@ -325,10 +321,8 @@ class Array[Shape: _Shape = _Shape]:
     # `reshape` below: an argument list cannot be captured as a `Flag`.
     @overload
     def transpose(self, *axes: int) -> Array[IntTuple]: ...
-    # Positional-only: JAX's signature is `reshape(self, *args, order=...)`, so
-    # there is no `shape` keyword to bind.
     @overload
-    def reshape[NewShape: Flag[_NewShape]](
+    def reshape[NewShape: _Shape](
         self,
         shape: NewShape,
         /,
@@ -337,15 +331,12 @@ class Array[Shape: _Shape = _Shape]:
         out_sharding: _NamedSharding | _PartitionSpec | None = ...,
     ) -> Array[reshape_shape(Shape, NewShape)]: ...
     @overload
-    def reshape[NewShape: IntTuple](self, *shape: *NewShape) -> Array[NewShape]: ...
-    # JAX's variadic spelling is accepted but intentionally not modeled: an
-    # argument list cannot be captured as a `Flag`, so the shape is gradual and
-    # `reshape_shape` never runs, which leaves the `-1` and negative-size checks
-    # to runtime. Accepting it unvalidated is the intended behavior rather than a
-    # gap to close: rejecting it would flag valid and very common JAX, and
-    # modeling it needs a Pyrefly capability that does not exist. The tuple
-    # spelling, `a.reshape((2, 6))`, is both exact and validated. Only the method
-    # is variadic; `jnp.reshape(a, 2, 6)` is an error in JAX itself.
+    def reshape[NewShape: _Shape](
+        self,
+        *shape: *NewShape,
+        order: str = ...,
+        out_sharding: _NamedSharding | _PartitionSpec | None = ...,
+    ) -> Array[reshape_shape(Shape, NewShape)]: ...
     @overload
     def reshape(
         self,

tensor-shapes/jax-pyrefly-stubs/jax-stubs/_shapes.pyi

@@ -382,21 +382,10 @@ def lax_reshape_shape(
     return new_sizes
 
 @type_shape_dsl_function
-def reshape_shape(shape: IntTuple, newshape: int | tuple[int, ...] | None) -> IntTuple:
-    # `None` is not a legal argument to `reshape`. The arm exists because an
-    # `int | tuple[int, ...]` parameter cannot be iterated after narrowing with
-    # `is_int_value` alone -- the DSL function silently evaluates to `Unknown`.
-    # Leading with the `None` check is what makes the narrowing work. The Torch
-    # stubs use the same workaround in `conv_shape`.
-    if newshape is None:
-        return dsl.Invalid("reshape requires a shape")
-    if dsl.is_int_value(newshape):
-        dims = (newshape,)
-    else:
-        dims = newshape
+def reshape_shape(shape: IntTuple, dims: IntTuple) -> IntTuple:
     if any(dsl.is_concrete_int(dim) and dim < -1 for dim in dims):
         return dsl.Invalid("reshape sizes must be -1 or non-negative")
-    inferred = tuple(dim for dim in dims if dsl.is_concrete_int(dim) and dim == -1)
+    inferred = tuple(1 for dim in dims if dsl.is_concrete_int(dim) and dim == -1)
     if len(inferred) > 1:
         return dsl.Invalid("reshape accepts at most one -1")
     known_shape = dsl.IntTuple(
@@ -407,7 +396,7 @@ def reshape_shape(shape: IntTuple, newshape: int | tuple[int, ...] | None) -> In
     if len(inferred) == 0:
         if dsl.is_concrete_int(total) and dsl.is_concrete_int(known) and total != known:
             return dsl.Invalid("reshape target element count does not match the input")
-        return dsl.IntTuple(dim for dim in dims)
+        return dims
     if dsl.is_concrete_int(known):
         if known == 0:
             return dsl.Invalid("could not infer size for dimension -1")
@@ -420,6 +409,11 @@ def reshape_shape(shape: IntTuple, newshape: int | tuple[int, ...] | None) -> In
         )
     )
 
+@type_shape_dsl_function
+def reshape_scalar_shape(shape: IntTuple, dim: Int) -> IntTuple:
+    dims = dsl.IntTuple((dim,))
+    return reshape_shape(shape, dims)
+
 @type_shape_dsl_function
 def tile_shape(shape: IntTuple, repeats: int | tuple[int, ...] | None) -> IntTuple:
     if repeats is None:

tensor-shapes/jax-pyrefly-stubs/jax-stubs/numpy/__init__.pyi

@@ -69,6 +69,7 @@ from jax._shapes import (
     ravel_shape,
     reduce_shape,
     repeat_shape,
+    reshape_scalar_shape,
     reshape_shape,
     reverse_shape,
     roll_shape,
@@ -146,9 +147,9 @@ from . import fft as fft, linalg as linalg
 type _Shape = IntTuple
 type _Axis = int | tuple[int, ...] | None
 type _SingleAxis = int | None
-# The trailing `None` is not a legal argument to `reshape`. It is present because
+# The trailing `None` is not a legal argument to `broadcast_to`. It is present because
 # an `int | tuple[int, ...]` parameter cannot be iterated inside a DSL function
-# after narrowing with `is_int_value` alone. See `reshape_shape`, which rejects it.
+# after narrowing with `is_int_value` alone. See `broadcast_to_shape`, which rejects it.
 type _NewShape = int | tuple[int, ...] | None
 type _Scalar = bool | int | float | complex | np.number
 type _ShapedArrayLike[Shape: _Shape] = _Array[Shape] | np.ndarray[Shape]
@@ -1209,7 +1210,7 @@ def matrix_transpose[Batch: IntTuple, M: IntVar, N: IntVar](
 # A single int or tuple, matching JAX: the free function is not variadic, so
 # `jnp.reshape(a, 2, 3)` is an error there. `Array.reshape` is the variadic one.
 @overload
-def reshape[NewShape: Flag[_NewShape], Shape: _Shape = []](
+def reshape[NewShape: _Shape, Shape: _Shape = []](
     a: _ArrayLike[Shape],
     shape: NewShape,
     order: str = ...,
@@ -1218,14 +1219,14 @@ def reshape[NewShape: Flag[_NewShape], Shape: _Shape = []](
     out_sharding: _NamedSharding | _PartitionSpec | None = ...,
 ) -> _Array[reshape_shape(Shape, NewShape)]: ...
 @overload
-def reshape[NewShape: _Shape](
-    a: _ArrayLike[Any],
-    shape: NewShape,
+def reshape[NewDim: Int, Shape: _Shape = []](
+    a: _ArrayLike[Shape],
+    shape: NewDim,
     order: str = ...,
     *,
     copy: bool | None = ...,
     out_sharding: _NamedSharding | _PartitionSpec | None = ...,
-) -> _Array[NewShape]: ...
+) -> _Array[reshape_scalar_shape(Shape, NewDim)]: ...
 @overload
 def reshape(
     a: _ArrayLike[Any],

tensor-shapes/jax-pyrefly-stubs/test/test_reshape.py

@@ -52,6 +52,9 @@ def test_reshape_accepts_the_variadic_method_spelling() -> None:
     # Only the method is variadic; `jnp.reshape(a, 2, 6)` is an error in JAX itself.
     assert_shape(a.reshape(2, 6).shape, (2, 6))
     assert_shape(a.reshape(2, 2, 3).shape, (2, 2, 3))
+    assert_shape(a.reshape(2, -1).shape, (2, 6))
+    assert_shape(a.reshape(1, 1, -1).shape, (1, 1, 12))
+    assert_shape((a.reshape(1, 1, -1) * jnp.ones((1, 2, 12))).shape, (1, 2, 12))
 
 
 def test_reshape_accepts_keywords_where_jax_does() -> None: