Track literal shapes in array constructors (#4874) - #4874
Open
stroxler wants to merge 5 commits into
Open
Conversation
Summary:
Infer `IntTuple` parameters shared by multiple union alternatives from every matching input member before committing a solution. Shape candidates use the existing gradual shape policy: preserve equal dimensions, widen differing dimensions at the same rank, and fall back to gradual `IntTuple` when ranks differ.
For example:
```python
type ArrayLike[Shape: IntTuple] = Array[Shape] | NdArray[Shape]
def as_array[Shape: IntTuple](value: ArrayLike[Shape]) -> Array[Shape]: ...
def check(value: Array[[2, 3]] | NdArray[[4, 3]]) -> None:
assert_type(as_array(value), Array[[Any, 3]])
```
The analysis probes compatible pairings independently, obtaining `[2, 3]` from the `Array` arms and `[4, 3]` from the `NdArray` arms. Their common shape is `[Any, 3]`: the differing first dimension widens while the shared second dimension remains precise. Ordinary alternatives such as `None` or `str` count as successful matches without contributing shape candidates, and the complete union is revalidated with the joined shape.
This is intentionally limited to unions whose inferred variables are all `IntTuple`-bounded. Ordinary generic inference is unchanged. Speculative branch checks restore solver and subset state so one union arm cannot contaminate another.
Differential Revision: D119124734
Summary: Introduce `shape_extensions.Scalar` as a first-class type family that binds scalar operands to rank zero, normalizes positive-rank specializations to `Never`, and preserves suspended shape relationships through generic calls and ordinary type operations. Keep the existing `ScalarAsShape` behavior isolated for compatibility. For example: ```python type ArrayLike[Shape: IntTuple] = Array[Shape] | Scalar[Shape] def array_like[Shape: IntTuple](x: ArrayLike[Shape]) -> Array[Shape]: ... assert_type(array_like(1), Array[[]]) impossible: Scalar[[2], int] = 1 # rejected: Scalar[[2], int] normalizes to Never ``` Differential Revision: D119041651
Summary:
The preceding diff gives unions of shape-indexed structural types a gradual join for a shared shape parameter. This diff validates that behavior through `Scalar` aliases and rejects redundant `Scalar` arms that would otherwise become the sole binder for an observable shape variable.
For example:
```python
type ArrayLike[Shape: IntTuple] = Array[Shape] | Scalar[Shape]
def normalize[Shape: IntTuple](x: ArrayLike[Shape]) -> Array[Shape]: ...
def use(x: Array[[2, 3]] | Array[[4, 3]]) -> None:
assert_type(normalize(x), Array[[Any, 3]])
```
The tests also cover reversed union order, scalar/rank-zero combinations, bounded unions, and signatures where the shape variable appears outside the projected union and therefore must not be widened.
Differential Revision: D119041652
Summary: Introduce a first-class shape-indexed `ArrayCoercible` type and contextually project plain rectangular scalar list literals into exact shapes. Keep unsupported containers and unrelated union alternatives on ordinary typing paths, preserve gradual tails for unknown elements, and allow zero-sized dimensions. For example: ```python def consume[Shape: IntTuple](x: ArrayCoercible[Shape, int]) -> Array[Shape]: ... assert_type(consume(1), Array[[]]) assert_type(consume([[1, 2], [3, 4]]), Array[[2, 2]]) assert_type(consume([[], []]), Array[[2, 0]]) consume([[1], [2, 3]]) # rejected as non-rectangular ``` Differential Revision: D119041653
Contributor
|
@stroxler has exported this pull request. If you are a Meta employee, you can view the originating Diff in D119041654. |
stroxler
changed the title
Summary: Pull Request resolved: #4874 Adopt `ArrayCoercible` in the NumPy, JAX, and PyTorch array constructors so ordinary nested scalar literals infer precise shapes. Preserve existing shaped inputs, integer-size `torch.Tensor` construction, and broad gradual fallbacks for unsupported or dynamic inputs. For example: ```python assert_type(np.array([[1, 2], [3, 4]]), np.ndarray[[2, 2], Any]) assert_type(jnp.asarray([[], []]), jnp.Array[[2, 0]]) assert_type(torch.tensor([1, 2, 3]), Tensor[[3]]) assert_type(torch.Tensor(2, 3), Tensor[[2, 3]]) # existing size constructor ``` Differential Revision: D119041654
meta-codesync
Bot
changed the title
meta-codesync
Bot
force-pushed
the
export-D119041654
branch
from
September 9, 2026 23:38
1ca5116 to
5411390
Compare
|
According to mypy_primer, this change doesn't affect type check results on a corpus of open source code. ✅ |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary:
Adopt
ArrayCoerciblein the NumPy, JAX, and PyTorch array constructors so ordinary nested scalar literals infer precise shapes. Preserve existing shaped inputs, integer-sizetorch.Tensorconstruction, and broad gradual fallbacks for unsupported or dynamic inputs.For example:
Differential Revision: D119041654