[QEDady]

In @type_shape_dsl_function, assigning an integer expression in one branch and a bare or narrowed integer parameter in another previously failed domain merging, forcing stub authors to write + 0 workarounds:

if axis < 0:
    norm_axis = axis + output_rank
else:
    norm_axis = axis  # previously required `axis + 0`

Because arithmetic produced DeferredInteger (or Dimension) while bare/narrowed parameters produced UnknownParameters (or a parameter-only ValueSet), merge_flows rejected the join, and conditional expressions (a if cond else b) called DslStaticKind::join directly without DeferredInteger handling.

High-level summary (in suggested review order)

1. DslStaticKind::integer_parameter_uses extracts TypeShapeDslParameterUse entries from bare and narrowed non-None parameter slots, and TypeShapeDslInputDomain::integer_domain / shared_integer_domain classify whether all contributing parameters resolve to the Dimension (Int, or Int | None narrowed to exclude None) or Flag (Flag[int], or a narrowable Flag union narrowed via dsl.is_int_value) integer domain.
2. DslValidator::join_value_kinds unifies value-domain joining for both merge_flows (if/elif/else) and validate_assignment_value (a if cond else b), joining Dimension and DeferredInteger values with compatible integer parameters while retaining parameter_uses alongside non_parameter_kinds: FLAG_INT on merged Flag locals so they remain input-only and cannot be returned or used directly as IntTuple elements.
3. validate_axis_reduce resolves DeferredInteger axis slots to DslIntegerDomain::Flag (matching validate_flag_slot), and classify_deferred_integer_expression keeps chained arithmetic on merged Int and Flag[int] locals domain-polymorphic.
4. Added positive and negative tests covering Flag[int], narrowed Flag unions, Int, narrowed Int | None, dsl.is_concrete_int, elif chains, conditional expressions, dsl.axis_reduce, helper forwarding, chained arithmetic, and invalid cross-domain / escape cases, and removed + 0 from the expand_dims_shape documentation example.

• I am a human, and I fully understand the code I am submitting

The code was written by AI, but I reviewed it to the best I can.

Follow-up

The existing + 0 workarounds in the stubs were left untouched to keep this PR focused, but it was verified locally that removing all 77 redundant + 0 sites across jax-pyrefly-stubs, pyrefly-torch-stubs, and pyrefly-numpy-stubs passes the stub test suites. I will send a follow-up PR to remove them once this lands.

Test Plan

Ran python3 test.py, plus tensor-shapes/run_tests.py --static-only on a local copy of the stubs with the redundant + 0 workarounds removed.

[github-actions]

✅ This PR meets all requirements from our pull request checklist.

Optional recommendations:

• If this PR addresses an open issue, please link the issue in the PR description and claim it.
• This PR is on the larger side. If parts of it could be merged on their own, consider splitting it into separate PRs for faster review. If you've already split it as far as is practical, no action is needed.

Note: this comment updates automatically after changes to the PR. Updates may take up to half an hour.

[meta-codesync]

This pull request has been imported. If you are a Meta employee, you can view this in D130935138. (Because this pull request was imported automatically, there will not be any future comments.)

[github-actions]

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

[stroxler]

This looks pretty good to me, thanks!

An AI review flagged that the integer_domain helper probably needs some deduplication.

The other potential concern is handling of code like this:

import shape_extensions.dsl as dsl
from shape_extensions import Int, IntTuple, type_shape_dsl_function

@type_shape_dsl_function
def test0(shape: IntTuple, d: Int, e: int, c1: bool, c2: bool) -> IntTuple:
    if c1:
        out = e + 1
    elif c2:
        out = e
    else:
        out = d
    return shape


@type_shape_dsl_function
def test1(shape: IntTuple, d: Int, e: int, c1: bool, c2: bool) -> IntTuple:
    if not c1 and not c2:
        out = d
    elif c1:
        out = e + 1
    else:
        out = e
    return shape

These are equivalent in complexity, but it appears that we're accepting the first one (where a constant Flag[int] value then widens into a symbolic Int) whereas we are not accepting the second (where the order of branches is flipped).

I think it makes sense to see if we can fix this. It's also okay to skip it, and maybe just add a bug= case if for some reason fixing the problem proves hard

[stroxler]

AI review flags duplication plus inconsistencies here.

We have similar logic at pyrefly/lib/alt/type_level_dsl.rs:75–79, 126–140, but the handling of contains(Int) is slightly different. We should probably consolidate this