From 3d147f9b54c7e70ef00391e28221e1c7f03ea0b5 Mon Sep 17 00:00:00 2001 From: Jake VanderPlas Date: Wed, 11 Nov 2026 12:26:26 -0800 Subject: [PATCH] [tensor-shapes] jax.numpy.reshape: support -1 dimensions --- .../jax-pyrefly-stubs/jax-stubs/_array.pyi | 23 ++++++------------- .../jax-pyrefly-stubs/jax-stubs/_shapes.pyi | 22 +++++++----------- .../jax-stubs/numpy/__init__.pyi | 15 ++++++------ .../jax-pyrefly-stubs/test/test_reshape.py | 3 +++ 4 files changed, 26 insertions(+), 37 deletions(-) diff --git a/tensor-shapes/jax-pyrefly-stubs/jax-stubs/_array.pyi b/tensor-shapes/jax-pyrefly-stubs/jax-stubs/_array.pyi index 0eb462f302..dab548741b 100644 --- a/tensor-shapes/jax-pyrefly-stubs/jax-stubs/_array.pyi +++ b/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, diff --git a/tensor-shapes/jax-pyrefly-stubs/jax-stubs/_shapes.pyi b/tensor-shapes/jax-pyrefly-stubs/jax-stubs/_shapes.pyi index aaa9358a12..517585513c 100644 --- a/tensor-shapes/jax-pyrefly-stubs/jax-stubs/_shapes.pyi +++ b/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: diff --git a/tensor-shapes/jax-pyrefly-stubs/jax-stubs/numpy/__init__.pyi b/tensor-shapes/jax-pyrefly-stubs/jax-stubs/numpy/__init__.pyi index 195f905538..af24f32bc8 100644 --- a/tensor-shapes/jax-pyrefly-stubs/jax-stubs/numpy/__init__.pyi +++ b/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], diff --git a/tensor-shapes/jax-pyrefly-stubs/test/test_reshape.py b/tensor-shapes/jax-pyrefly-stubs/test/test_reshape.py index e345fee2a6..cdfa3ecf06 100644 --- a/tensor-shapes/jax-pyrefly-stubs/test/test_reshape.py +++ b/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: