[WIP] Proposal: coercible shape_extensions.Scalar

[jakevdp]

DO NOT MERGE – draft only

This PR is a draft/proposal for a new shape_extensions.Scalar type that can be used to simplify APIs in NumPy and related libraries that allow scalars to be coerced into arrays.

Currently, and API like jax.numpy.add requires multiple overloads to allow mixing of scalar and array values:

from shape_extensions import IntTuple

type Scalar = bool | int | float | complex

@overload
def add[Shape: IntTuple](x1: Array[Shape], x2: Scalar) -> Array[Shape]: ...
@overload
def add[Shape: IntTuple](x1: Scalar, x2: Array[Shape]) -> Array[Shape]: ...
@overload
def add[Shape1: IntTuple, Shape2: IntTuple](
    x1: Array[Shape1],
    x2: Array[Shape2]
) -> Array[broadcast(Shape1, Shape2)]: ...

This PR contains a proposed shape_extensions.Scalar which allows this common pattern to be expressed much more succinctly:

from shape_extensions import IntTuple, Scalar

type ArrayLike[Shape: IntTuple] = Array[Shape] | Scalar[Shape]

def add[Shape1: IntTuple, Shape2: IntTuple](
    x1: ArrayLike[Shape1],
    x2: ArrayLike[Shape2]
) -> Array[broadcast(Shape1, Shape2)]: ...

The Scalar type defined here is specially recognized by pyrefly, and is only valid when its argument is an empty shape. This would allow libraries to compactly specify APIs which allow scalar inputs to be coerced to scalar arrays without the need for many overloaded signatures.

This PR adds a proposed Scalar implementation, modifies jax.numpy.add to make use of it, and adds a test to cover its intended semantics.

A few open questions:

• how to support library-defined scalar types, e.g. numpy.number? Perhaps the Scalar API should be something more general, like Scalar[Shape, int | float | complex]?
• how to support dtype annotations like those used by NumPy?
• Scalar[Shape] would admit arbitrary shapes in Python's general type system. Can we solve this a different way?
• other considerations I'm not thinking of?

[meta-codesync]

This pull request has been imported. If you are a Meta employee, you can view this in D118713401. (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. ✅