[WIP] Proposal: coercible shape_extensions.Scalar

pyrefly/lib/alt/expr.rs

@@ -5133,8 +5133,24 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> {
         {
             return Type::ClassType(base_class);
         }
-        self.shaped_array_classtype_to_shaped_array_type(&base_class)
-            .to_type()
+        let shaped_array = self.shaped_array_classtype_to_shaped_array_type(&base_class);
+        if cls.has_toplevel_qname("shape_extensions", "Scalar") {
+            let is_valid = match shaped_array.shape().view() {
+                IntTupleView::Concrete(dims) => dims.is_empty(),
+                IntTupleView::Unpacked { .. } => true,
+                IntTupleView::Gradual => false,
+            };
+            if !is_valid {
+                self.error(
+                    errors,
+                    range,
+                    ErrorKind::InvalidAnnotation,
+                    "`shape_extensions.Scalar` only accepts an empty shape `[]`".to_owned(),
+                );
+                return Type::any_error();
+            }
+        }
+        shaped_array.to_type()
     }
 
     pub(crate) fn parse_int_tuple_type(

pyrefly/lib/alt/solve.rs

@@ -6603,7 +6603,14 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> {
                             // Canonicalize bare shaped-array types to Type::ShapedArray(shapeless)
                             // for consistency. Subscripted arrays are already converted to
                             // Type::ShapedArray during annotation parsing, so only the bare case reaches here.
-                            Some(ShapedArrayType::shapeless(cls.clone()).to_type())
+                            if cls.has_qname("shape_extensions", "Scalar") {
+                                Some(
+                                    ShapedArrayType::new(cls.clone(), IntTuple::new(Vec::new()))
+                                        .to_type(),
+                                )
+                            } else {
+                                Some(ShapedArrayType::shapeless(cls.clone()).to_type())
+                            }
                         } else if cls.has_qname("types", "NoneType") {
                             // Normalize type[NoneType] as None
                             Some(self.heap.mk_none())

pyrefly/lib/solver/subset.rs

@@ -142,6 +142,20 @@ fn is_int_class_type(cls: &ClassType) -> bool {
     cls.has_qname("shape_extensions", "Int")
 }
 
+fn is_scalar_type(ty: &Type) -> bool {
+    match ty {
+        Type::ClassType(cls) => {
+            cls.is_builtin("int")
+                || cls.is_builtin("float")
+                || cls.is_builtin("complex")
+                || cls.is_builtin("bool")
+        }
+        Type::Literal(lit) => matches!(lit.value, Lit::Int(_) | Lit::Bool(_)),
+        Type::Int(_) => true,
+        _ => false,
+    }
+}
+
 fn params_have_any_args_and_kwargs(params: &Params) -> bool {
     match params {
         Params::List(args) | Params::Partial(args) => params_are_gradual_variadic(args.items()),
@@ -2282,6 +2296,24 @@ impl<'solver, 'subset, Ans: LookupAnswer> Subset<'solver, 'subset, Ans> {
             (Type::ShapedArray(got_shaped_array), Type::ShapedArray(want_shaped_array)) => {
                 self.is_subset_shaped_array(got_shaped_array, want_shaped_array)
             }
+            (got, Type::ShapedArray(want_shaped_array))
+                if is_scalar_type(got)
+                    && want_shaped_array
+                        .base_class
+                        .has_qname("shape_extensions", "Scalar") =>
+            {
+                let scalar_shape = IntTuple::new(Vec::new());
+                let scalar_shape_arg = scalar_shape.to_shape_arg_type();
+                let (_, want_arg) = self.shape_param_and_arg(want_shaped_array)?;
+                if IntTuple::from_shape_arg_type(want_arg)
+                    .or_else(|| tuple_carrier_to_shape(want_arg))
+                    .is_none()
+                {
+                    self.is_subset_eq(&scalar_shape_arg, want_arg)
+                } else {
+                    self.bind_tensor_dimensions(&scalar_shape, &want_shaped_array.shape())
+                }
+            }
             // Tensor is subtype of its base class
             (Type::ShapedArray(tensor), Type::ClassType(cls)) => {
                 let got = self.shaped_array_as_carrier_class(tensor)?;

tensor-shapes/pyrefly-jax-stubs/jax-stubs/numpy/__init__.pyi

@@ -49,17 +49,21 @@ from shape_extensions import (
     IntTuples,
     IntVar,
     MapIntTuples,
+    Scalar,
 )
 
 from . import fft as fft, linalg as linalg
 
 type _Shape = IntTuple
+type _AnyShape = tuple[Any, ...]
 type _Axis = int | tuple[int, ...] | None
 # 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
 
+type ArrayLike[Shape: _Shape = _AnyShape] = Array[Shape] | Scalar[Shape]
+
 # Ranks 1 through 3 are exact; any other integer sequence, including a longer
 # tuple or a list, falls through to a gradual overload rather than being
 # rejected.
@@ -657,17 +661,8 @@ def unwrap[Shape: _Shape](
 
 # Broadcasting elementwise binary functions. Each takes a scalar in either
 # position as well as an array: rejecting `jnp.add(a, 1)` would flag valid code.
-@overload
-def add[Shape: _Shape](
-    x1: Array[Shape], x2: int | float | complex, /
-) -> Array[Shape]: ...
-@overload
-def add[Shape: _Shape](
-    x1: int | float | complex, x2: Array[Shape], /
-) -> Array[Shape]: ...
-@overload
 def add[Shape1: _Shape, Shape2: _Shape](
-    x1: Array[Shape1], x2: Array[Shape2], /
+    x1: ArrayLike[Shape1], x2: ArrayLike[Shape2], /
 ) -> Array[broadcast(Shape1, Shape2)]: ...
 @overload
 def arctan2[Shape: _Shape](

tensor-shapes/pyrefly-jax-stubs/test/test_arithmetic.py

@@ -6,7 +6,7 @@
 from __future__ import annotations
 
 import jax.numpy as jnp
-from shape_extensions import assert_shape
+from shape_extensions import assert_shape, Scalar
 
 
 def test_elementwise_operators_preserve_shape() -> None:
@@ -236,3 +236,48 @@ def test_binary_functions_reject_incompatible_broadcast() -> None:
         pass
     else:
         raise AssertionError("expected JAX to reject incompatible shapes")
+
+
+def test_add_arraylike_annotations() -> None:
+    # Scalar / scalar operations produce 0-D arrays.
+    assert_shape(jnp.add(1, 2), ())
+    assert_shape(jnp.add(1.5, 2.5), ())
+    assert_shape(jnp.add(1, 2.0), ())
+    assert_shape(jnp.add(True, 1), ())
+    assert_shape(jnp.add(True, False), ())
+    assert_shape(jnp.add(1j, 2.0), ())
+    assert_shape(jnp.add(1j, 2j), ())
+
+    # Scalar / array operations preserve or broadcast shape.
+    scalar_0d = jnp.ones(())
+    assert_shape(jnp.add(1, scalar_0d), ())
+    assert_shape(jnp.add(scalar_0d, 1), ())
+    assert_shape(jnp.add(scalar_0d, scalar_0d), ())
+
+    vector = jnp.ones(4)
+    assert_shape(jnp.add(vector, 1), (4,))
+    assert_shape(jnp.add(1, vector), (4,))
+    assert_shape(jnp.add(vector, 2.5), (4,))
+    assert_shape(jnp.add(2.5, vector), (4,))
+
+    matrix = jnp.ones((3, 4))
+    assert_shape(jnp.add(matrix, 1), (3, 4))
+    assert_shape(jnp.add(1, matrix), (3, 4))
+    assert_shape(jnp.add(matrix, 2.5), (3, 4))
+    assert_shape(jnp.add(2.5, matrix), (3, 4))
+    assert_shape(jnp.add(matrix, 1j), (3, 4))
+    assert_shape(jnp.add(1j, matrix), (3, 4))
+    assert_shape(jnp.add(matrix, True), (3, 4))
+    assert_shape(jnp.add(True, matrix), (3, 4))
+
+    tensor = jnp.ones((2, 3, 4))
+    assert_shape(jnp.add(tensor, 1), (2, 3, 4))
+    assert_shape(jnp.add(1, tensor), (2, 3, 4))
+    assert_shape(jnp.add(tensor, 1.0), (2, 3, 4))
+    assert_shape(jnp.add(1.0, tensor), (2, 3, 4))
+
+    # Scalar type only allows empty shapes:
+    _: Scalar[()] = 1
+    _bare: Scalar = 1
+    # E: `shape_extensions.Scalar` only accepts an empty shape `[]`
+    _bad: Scalar[[3, 4]]

tensor-shapes/pyrefly-shape-extensions/shape_extensions/__init__.py

@@ -24,6 +24,7 @@
     "IntVar",
     "MapIntTuples",
     "ProxyMethod",
+    "Scalar",
     "SymbolicArithExpr",
     "TypeVarTuple",
     "assert_shape",
@@ -279,6 +280,13 @@ def decorator(cls: type) -> type:
     return decorator
 
 
+@shaped_array(shape="Shape")
+class Scalar[Shape: IntTuple = tuple[()]]:
+    """Marker type for scalar types that coerce to an empty shape `[]`."""
+
+    pass
+
+
 class MapIntTuples:
     """Map a unary type lambda over an ``IntTuples`` value.