Skip to content

Add shorthand and test IntTuple tparam defaults (#4896) - #4896

Open
stroxler wants to merge 1 commit into
mainfrom
export-D119566820
Open

Add shorthand and test IntTuple tparam defaults (#4896)#4896
stroxler wants to merge 1 commit into
mainfrom
export-D119566820

Conversation

@stroxler

@stroxler stroxler commented Sep 10, 2026

Copy link
Copy Markdown
Contributor

Summary:

In jax and other array libraries, it's not unusual to auto-promote scalars to empty-shape
arrays so that broadcasting works. For example, we can have

def add[S0: IntTuple, S1: IntTuple](
    x: Array[S0] | float, y: Array[S1] | float
) -> Array[broadcast(S0, S1)]

and we want something like f(3.141, y) to produce the same shape as y.

The problem is that this doesn't actually work: if we pass a float for x, we wind up failing
to constrain S0, and it solves to gradual IntTuple, which means we drop the shape.

The workaround is to use a default:

def add[S0: IntTuple = [], S1: IntTuple = []](
    x: Array[S0] | float, y: Array[S1] | float
) -> Array[broadcast(S0, S1)]

so that taking the scalar branch of the union leads to the unconstrained shape parameter being solved to [] by the default.

This diff adds support for the [] shorthand in this syntactic position for IntTuple-constrained type parameters, and verifies that it works as expected

Differential Revision: D119566820

@meta-cla meta-cla Bot added the cla signed label Sep 10, 2026
@meta-codesync

meta-codesync Bot commented Sep 10, 2026

Copy link
Copy Markdown
Contributor

@stroxler has exported this pull request. If you are a Meta employee, you can view the originating Diff in D119566820.

@github-actions

This comment has been minimized.

meta-codesync Bot pushed a commit that referenced this pull request Sep 11, 2026
Summary:

In jax and other array libraries, it's not unusual to auto-promote scalars to empty-shape
arrays so that broadcasting works. For example, we can have
```
def add[S0: IntTuple, S1: IntTuple](
    x: Array[S0] | float, y: Array[S1] | float
) -> Array[broadcast(S0, S1)]
```
and we want something like `f(3.141, y)` to produce the same shape as `y`.

The problem is that this doesn't actually work: if we pass a `float` for `x`, we wind up failing
to constrain `S0`, and it solves to gradual `IntTuple`, which means we drop the shape.

The workaround is to use a default:
```
def add[S0: IntTuple = [], S1: IntTuple = []](
    x: Array[S0] | float, y: Array[S1] | float
) -> Array[broadcast(S0, S1)]
```
so that taking the scalar branch of the union leads to the unconstrained shape parameter being solved to `[]` by the default.

This diff adds support for the `[]` shorthand in this syntactic position for `IntTuple`-constrained type parameters, and verifies that it works as expected

Differential Revision: D119566820
meta-codesync Bot pushed a commit that referenced this pull request Sep 11, 2026
Summary:

In jax and other array libraries, it's not unusual to auto-promote scalars to empty-shape
arrays so that broadcasting works. For example, we can have
```
def add[S0: IntTuple, S1: IntTuple](
    x: Array[S0] | float, y: Array[S1] | float
) -> Array[broadcast(S0, S1)]
```
and we want something like `f(3.141, y)` to produce the same shape as `y`.

The problem is that this doesn't actually work: if we pass a `float` for `x`, we wind up failing
to constrain `S0`, and it solves to gradual `IntTuple`, which means we drop the shape.

The workaround is to use a default:
```
def add[S0: IntTuple = [], S1: IntTuple = []](
    x: Array[S0] | float, y: Array[S1] | float
) -> Array[broadcast(S0, S1)]
```
so that taking the scalar branch of the union leads to the unconstrained shape parameter being solved to `[]` by the default.

This diff adds support for the `[]` shorthand in this syntactic position for `IntTuple`-constrained type parameters, and verifies that it works as expected

Differential Revision: D119566820
meta-codesync Bot pushed a commit that referenced this pull request Sep 11, 2026
Summary:

In jax and other array libraries, it's not unusual to auto-promote scalars to empty-shape
arrays so that broadcasting works. For example, we can have
```
def add[S0: IntTuple, S1: IntTuple](
    x: Array[S0] | float, y: Array[S1] | float
) -> Array[broadcast(S0, S1)]
```
and we want something like `f(3.141, y)` to produce the same shape as `y`.

The problem is that this doesn't actually work: if we pass a `float` for `x`, we wind up failing
to constrain `S0`, and it solves to gradual `IntTuple`, which means we drop the shape.

The workaround is to use a default:
```
def add[S0: IntTuple = [], S1: IntTuple = []](
    x: Array[S0] | float, y: Array[S1] | float
) -> Array[broadcast(S0, S1)]
```
so that taking the scalar branch of the union leads to the unconstrained shape parameter being solved to `[]` by the default.

This diff adds support for the `[]` shorthand in this syntactic position for `IntTuple`-constrained type parameters, and verifies that it works as expected

Differential Revision: D119566820
meta-codesync Bot pushed a commit that referenced this pull request Sep 11, 2026
Summary:

In jax and other array libraries, it's not unusual to auto-promote scalars to empty-shape
arrays so that broadcasting works. For example, we can have
```
def add[S0: IntTuple, S1: IntTuple](
    x: Array[S0] | float, y: Array[S1] | float
) -> Array[broadcast(S0, S1)]
```
and we want something like `f(3.141, y)` to produce the same shape as `y`.

The problem is that this doesn't actually work: if we pass a `float` for `x`, we wind up failing
to constrain `S0`, and it solves to gradual `IntTuple`, which means we drop the shape.

The workaround is to use a default:
```
def add[S0: IntTuple = [], S1: IntTuple = []](
    x: Array[S0] | float, y: Array[S1] | float
) -> Array[broadcast(S0, S1)]
```
so that taking the scalar branch of the union leads to the unconstrained shape parameter being solved to `[]` by the default.

This diff adds support for the `[]` shorthand in this syntactic position for `IntTuple`-constrained type parameters, and verifies that it works as expected

Differential Revision: D119566820
Summary:

In jax and other array libraries, it's not unusual to auto-promote scalars to empty-shape
arrays so that broadcasting works. For example, we can have
```
def add[S0: IntTuple, S1: IntTuple](
    x: Array[S0] | float, y: Array[S1] | float
) -> Array[broadcast(S0, S1)]
```
and we want something like `f(3.141, y)` to produce the same shape as `y`.

The problem is that this doesn't actually work: if we pass a `float` for `x`, we wind up failing
to constrain `S0`, and it solves to gradual `IntTuple`, which means we drop the shape.

The workaround is to use a default:
```
def add[S0: IntTuple = [], S1: IntTuple = []](
    x: Array[S0] | float, y: Array[S1] | float
) -> Array[broadcast(S0, S1)]
```
so that taking the scalar branch of the union leads to the unconstrained shape parameter being solved to `[]` by the default.

This diff adds support for the `[]` shorthand in this syntactic position for `IntTuple`-constrained type parameters, and verifies that it works as expected

Differential Revision: D119566820
@meta-codesync meta-codesync Bot changed the title Add shorthand and test IntTuple tparam defaults Add shorthand and test IntTuple tparam defaults (#4896) Sep 11, 2026
@github-actions github-actions Bot added size/m and removed size/m labels Sep 11, 2026
meta-codesync Bot pushed a commit that referenced this pull request Sep 11, 2026
Summary:

In jax and other array libraries, it's not unusual to auto-promote scalars to empty-shape
arrays so that broadcasting works. For example, we can have
```
def add[S0: IntTuple, S1: IntTuple](
    x: Array[S0] | float, y: Array[S1] | float
) -> Array[broadcast(S0, S1)]
```
and we want something like `f(3.141, y)` to produce the same shape as `y`.

The problem is that this doesn't actually work: if we pass a `float` for `x`, we wind up failing
to constrain `S0`, and it solves to gradual `IntTuple`, which means we drop the shape.

The workaround is to use a default:
```
def add[S0: IntTuple = [], S1: IntTuple = []](
    x: Array[S0] | float, y: Array[S1] | float
) -> Array[broadcast(S0, S1)]
```
so that taking the scalar branch of the union leads to the unconstrained shape parameter being solved to `[]` by the default.

This diff adds support for the `[]` shorthand in this syntactic position for `IntTuple`-constrained type parameters, and verifies that it works as expected

Differential Revision: D119566820
@github-actions

Copy link
Copy Markdown

According to mypy_primer, this change doesn't affect type check results on a corpus of open source code. ✅

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant