[tensor-shapes] update for JAX v0.12.0

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

@@ -106,7 +106,6 @@ from jax import (
 from jax._src.ad_checkpoint import (
     checkpoint as checkpoint,
     checkpoint_policies as checkpoint_policies,
-    custom_remat as custom_remat,
     remat as remat,
 )
 from jax._src.api import (

tensor-shapes/jax-pyrefly-stubs/jax-stubs/_shapes.pyi

@@ -124,26 +124,30 @@ def parameter_broadcast_shape(
     return result_shape
 
 @type_shape_dsl_function
-def multinomial_shape(
-    n_shape: IntTuple, p_shape: IntTuple, result_shape: IntTuple
+def multinomial_shape(n_shape: IntTuple, p_shape: IntTuple) -> IntTuple:
+    batch_shape = p_shape[:-1]
+    return multinomial_sample_shape(n_shape, p_shape, batch_shape)
+
+@type_shape_dsl_function
+def multinomial_sample_shape(
+    n_shape: IntTuple, p_shape: IntTuple, sample_shape: IntTuple
 ) -> IntTuple:
-    if len(p_shape) == 0 or len(result_shape) == 0:
+    if len(p_shape) == 0:
         return dsl.Invalid("multinomial probabilities must have an event dimension")
-    probability_shapes = dsl.IntTuples((p_shape, result_shape))
+    probability_shapes = dsl.IntTuples((p_shape[:-1], sample_shape))
     probability_shape = gufunc_broadcast("(),()->()", probability_shapes)
-    if len(probability_shape) != len(result_shape) or any(
-        probability_shape[index] != result_shape[index]
-        for index in range(len(result_shape))
+    if len(probability_shape) != len(sample_shape) or any(
+        probability_shape[index] != sample_shape[index]
+        for index in range(len(sample_shape))
     ):
         return dsl.Invalid("probabilities cannot broadcast to the requested shape")
-    count_shapes = dsl.IntTuples((n_shape, result_shape[:-1]))
+    count_shapes = dsl.IntTuples((n_shape, sample_shape))
     count_shape = gufunc_broadcast("(),()->()", count_shapes)
-    if len(count_shape) != len(result_shape) - 1 or any(
-        count_shape[index] != result_shape[index]
-        for index in range(len(result_shape) - 1)
+    if len(count_shape) != len(sample_shape) or any(
+        count_shape[index] != sample_shape[index] for index in range(len(sample_shape))
     ):
         return dsl.Invalid("counts cannot broadcast to the requested batch shape")
-    return result_shape
+    return dsl.concat(sample_shape, p_shape[-1:])
 
 @type_shape_dsl_function
 def multivariate_normal_shape(

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

@@ -10,6 +10,7 @@ from typing import (
     Literal,
     NamedTuple,
     overload,
+    Protocol,
     Sequence,
     Unpack,
 )
@@ -1625,19 +1626,37 @@ def atleast_1d[Shape: _Shape = []](
     ary: _ArrayLike[Shape], /
 ) -> _Array[atleast_1d_shape(Shape)]: ...
 @overload
-def atleast_1d(*arys: _ArrayLike[Any]) -> list[_Array[IntTuple]]: ...
+def atleast_1d[Shapes: IntTuples](
+    *arys: Unpack[MapIntTuples[lambda S: _Array[S], Shapes]],
+) -> MapIntTuples[lambda S: _Array[atleast_1d_shape(S)], Shapes]: ...
+@overload
+def atleast_1d[Shapes: IntTuples](
+    *arys: Unpack[MapIntTuples[lambda S: _ArrayLike[S], Shapes]],
+) -> MapIntTuples[lambda S: _Array[S], Shapes]: ...
 @overload
 def atleast_2d[Shape: _Shape = []](
     ary: _ArrayLike[Shape], /
 ) -> _Array[atleast_2d_shape(Shape)]: ...
 @overload
-def atleast_2d(*arys: _ArrayLike[Any]) -> list[_Array[IntTuple]]: ...
+def atleast_2d[Shapes: IntTuples](
+    *arys: Unpack[MapIntTuples[lambda S: _Array[S], Shapes]],
+) -> MapIntTuples[lambda S: _Array[atleast_2d_shape(S)], Shapes]: ...
+@overload
+def atleast_2d[Shapes: IntTuples](
+    *arys: Unpack[MapIntTuples[lambda S: _ArrayLike[S], Shapes]],
+) -> MapIntTuples[lambda S: _Array[S], Shapes]: ...
 @overload
 def atleast_3d[Shape: _Shape = []](
     ary: _ArrayLike[Shape], /
 ) -> _Array[atleast_3d_shape(Shape)]: ...
 @overload
-def atleast_3d(*arys: _ArrayLike[Any]) -> list[_Array[IntTuple]]: ...
+def atleast_3d[Shapes: IntTuples](
+    *arys: Unpack[MapIntTuples[lambda S: _Array[S], Shapes]],
+) -> MapIntTuples[lambda S: _Array[atleast_3d_shape(S)], Shapes]: ...
+@overload
+def atleast_3d[Shapes: IntTuples](
+    *arys: Unpack[MapIntTuples[lambda S: _ArrayLike[S], Shapes]],
+) -> MapIntTuples[lambda S: _Array[S], Shapes]: ...
 
 # JAX accepts any integer sequence for an axis, but only a tuple is a Flag
 # domain, so any other sequence yields a gradual shape. Rejecting it would flag
@@ -3664,44 +3683,53 @@ def roots(p: _ArrayLike[Any], *, strip_zeros: bool = True) -> _Array[IntTuple]:
 
 # Scalar constructors
 
-bool: Any
-bool_: Any
-int_: Any
-int8: Any
-int16: Any
-int32: Any
-int64: Any
-uint: Any
-uint8: Any
-uint16: Any
-uint32: Any
-uint64: Any
-int1: Any
-int2: Any
-int4: Any
-uint1: Any
-uint2: Any
-uint4: Any
-float_: Any
-float16: Any
-float32: Any
-float64: Any
-bfloat16: Any
-single: Any
-double: Any
-csingle: Any
-cdouble: Any
-complex_: Any
-complex64: Any
-complex128: Any
-float4_e2m1fn: Any
-float6_e2m3fn: Any
-float6_e3m2fn: Any
-float8_e3m4: Any
-float8_e4m3: Any
-float8_e4m3b11fnuz: Any
-float8_e4m3fn: Any
-float8_e4m3fnuz: Any
-float8_e5m2: Any
-float8_e5m2fnuz: Any
-float8_e8m0fnu: Any
+class _ScalarType(Protocol):
+    @property
+    def dtype(self) -> np.dtype: ...
+    @property
+    def __numpy_dtype__(self) -> np.dtype: ...
+    def __call__(self, x: Any, /) -> _Array: ...
+    def __hash__(self, /) -> int: ...
+    def __eq__(self, other: Any, /) -> bool: ...
+
+bool: _ScalarType
+bool_: _ScalarType
+int_: _ScalarType
+int8: _ScalarType
+int16: _ScalarType
+int32: _ScalarType
+int64: _ScalarType
+uint: _ScalarType
+uint8: _ScalarType
+uint16: _ScalarType
+uint32: _ScalarType
+uint64: _ScalarType
+int1: _ScalarType
+int2: _ScalarType
+int4: _ScalarType
+uint1: _ScalarType
+uint2: _ScalarType
+uint4: _ScalarType
+float_: _ScalarType
+float16: _ScalarType
+float32: _ScalarType
+float64: _ScalarType
+bfloat16: _ScalarType
+single: _ScalarType
+double: _ScalarType
+csingle: _ScalarType
+cdouble: _ScalarType
+complex_: _ScalarType
+complex64: _ScalarType
+complex128: _ScalarType
+float4_e2m1fn: _ScalarType
+float6_e2m3fn: _ScalarType
+float6_e3m2fn: _ScalarType
+float8_e3m4: _ScalarType
+float8_e4m3: _ScalarType
+float8_e4m3b11fnuz: _ScalarType
+float8_e4m3fn: _ScalarType
+float8_e4m3fnuz: _ScalarType
+float8_e5m2: _ScalarType
+float8_e5m2fnuz: _ScalarType
+float8_e8m0fnu: _ScalarType

tensor-shapes/jax-pyrefly-stubs/jax-stubs/random.pyi

@@ -13,6 +13,7 @@ from jax._shapes import (
     double_sided_maxwell_shape,
     event_sample_shape,
     event_shape,
+    multinomial_sample_shape,
     multinomial_shape,
     multivariate_normal_shape,
     orthogonal_shape,
@@ -913,7 +914,7 @@ def multinomial[NShape: _Shape = [], PShape: _Shape = []](
     shape: None = None,
     dtype: DTypeLike | None = None,
     unroll: int | bool = 1,
-) -> _Array[multinomial_shape(NShape, PShape, PShape)]: ...
+) -> _Array[multinomial_shape(NShape, PShape)]: ...
 @overload
 def multinomial[NShape: _Shape = [], PShape: _Shape = [], Shape: _Shape = []](
     key: _Array[IntTuple],
@@ -923,7 +924,7 @@ def multinomial[NShape: _Shape = [], PShape: _Shape = [], Shape: _Shape = []](
     shape: Shape,
     dtype: DTypeLike | None = None,
     unroll: int | bool = 1,
-) -> _Array[multinomial_shape(NShape, PShape, Shape)]: ...
+) -> _Array[multinomial_sample_shape(NShape, PShape, Shape)]: ...
 @overload
 def multinomial(
     key: _Array[IntTuple],

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

@@ -165,8 +165,9 @@ def test_event_distributions() -> None:
     probabilities = jnp.full((2, 3), 1 / 3)
     assert_shape(random.multinomial(key, 5, probabilities).shape, (2, 3))
     assert_shape(
-        random.multinomial(key, 5, probabilities, shape=(4, 2, 3)).shape, (4, 2, 3)
+        random.multinomial(key, 5, probabilities, shape=(4, 2)).shape, (4, 2, 3)
     )
+    assert_shape(random.multinomial(key, 5, probabilities[0], shape=()).shape, (3,))
 
     try:
         # E: parameters cannot broadcast to the requested shape
@@ -187,7 +188,7 @@ def test_event_distributions() -> None:
     try:
         # E: multinomial probabilities must have an event dimension
         random.multinomial(key, 5, 1.0)
-    except ValueError:
+    except IndexError:
         pass
     else:
         raise AssertionError("expected JAX to reject scalar multinomial probabilities")
@@ -200,6 +201,16 @@ def test_event_distributions() -> None:
     else:
         raise AssertionError("expected JAX to reject incompatible multinomial counts")
 
+    try:
+        # E: probabilities cannot broadcast to the requested shape
+        random.multinomial(key, 5, probabilities, shape=(4, 1))
+    except ValueError:
+        pass
+    else:
+        raise AssertionError(
+            "expected JAX to reject incompatible multinomial batch shape"
+        )
+
 
 def test_multivariate_normal() -> None:
     key = random.key(0)

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

@@ -185,18 +185,32 @@ def test_atleast_1d() -> None:
     assert_shape(jnp.atleast_1d(jnp.ones((2, 3, 4))).shape, (2, 3, 4))
     assert_shape(jnp.atleast_1d(np.ones((2, 3))).shape, (2, 3))
 
-    # Zero arguments: returns list of arrays
+    # Zero arguments: returns empty tuple
     res0 = jnp.atleast_1d()
-    assert_type(res0, list[jax.Array])
-    assert res0 == []
+    assert_type(res0, tuple[()])
+    assert res0 == ()
 
-    # Two arguments: returns list of arrays
-    res2 = jnp.atleast_1d(1, jnp.ones((2, 3)))
-    assert_type(res2, list[jax.Array])
-    assert isinstance(res2, list) and len(res2) == 2
+    # Two arguments: returns 2-tuple of shape-typed arrays
+    res2 = jnp.atleast_1d(jnp.ones(()), jnp.ones((2, 3)))
+    assert isinstance(res2, tuple) and len(res2) == 2
     assert all(isinstance(x, jax.Array) for x in res2)
-    assert res2[0].shape == (1,)
-    assert res2[1].shape == (2, 3)
+    assert_shape(res2[0].shape, (1,))
+    assert_shape(res2[1].shape, (2, 3))
+
+    # Three array arguments: returns 3-tuple of shape-typed arrays
+    res3_arr = jnp.atleast_1d(jnp.ones(()), jnp.ones(4), jnp.ones((2, 3)))
+    assert_type(res3_arr, tuple[jax.Array[[1]], jax.Array[[4]], jax.Array[[2, 3]]])
+    assert_shape(res3_arr[0].shape, (1,))
+    assert_shape(res3_arr[1].shape, (4,))
+    assert_shape(res3_arr[2].shape, (2, 3))
+
+    # Three mixed arguments (with scalar): returns 3-tuple with gradual shapes
+    res3 = jnp.atleast_1d(1, jnp.ones(4), jnp.ones((2, 3)))
+    assert_type(res3, tuple[jax.Array, jax.Array, jax.Array])
+    assert isinstance(res3, tuple) and len(res3) == 3
+    assert_shape(res3[0].shape, IntTuple, runtime=(1,))
+    assert_shape(res3[1].shape, IntTuple, runtime=(4,))
+    assert_shape(res3[2].shape, IntTuple, runtime=(2, 3))
 
 
 def test_atleast_2d() -> None:
@@ -213,18 +227,35 @@ def test_atleast_2d() -> None:
     assert_shape(jnp.atleast_2d(jnp.ones((2, 3, 4))).shape, (2, 3, 4))
     assert_shape(jnp.atleast_2d(np.ones((2, 3))).shape, (2, 3))
 
-    # Zero arguments: returns list of arrays
+    # Zero arguments: returns empty tuple
     res0 = jnp.atleast_2d()
-    assert_type(res0, list[jax.Array])
-    assert res0 == []
+    assert_type(res0, tuple[()])
+    assert res0 == ()
 
-    # Two arguments: returns list of arrays
+    # Two arguments: returns 2-tuple of shape-typed arrays
     res2 = jnp.atleast_2d(jnp.ones(3), jnp.ones((2, 3)))
-    assert_type(res2, list[jax.Array])
-    assert isinstance(res2, list) and len(res2) == 2
+    assert isinstance(res2, tuple) and len(res2) == 2
     assert all(isinstance(x, jax.Array) for x in res2)
-    assert res2[0].shape == (1, 3)
-    assert res2[1].shape == (2, 3)
+    assert_shape(res2[0].shape, (1, 3))
+    assert_shape(res2[1].shape, (2, 3))
+
+    # Three array arguments: returns 3-tuple of shape-typed arrays
+    res3_arr = jnp.atleast_2d(jnp.ones(()), jnp.ones(3), jnp.ones((2, 3)))
+    assert_type(
+        res3_arr,
+        tuple[jax.Array[[1, 1]], jax.Array[[1, 3]], jax.Array[[2, 3]]],
+    )
+    assert_shape(res3_arr[0].shape, (1, 1))
+    assert_shape(res3_arr[1].shape, (1, 3))
+    assert_shape(res3_arr[2].shape, (2, 3))
+
+    # Three mixed arguments (with scalar): returns 3-tuple with gradual shapes
+    res3 = jnp.atleast_2d(1, jnp.ones(3), jnp.ones((2, 3)))
+    assert_type(res3, tuple[jax.Array, jax.Array, jax.Array])
+    assert isinstance(res3, tuple) and len(res3) == 3
+    assert_shape(res3[0].shape, IntTuple, runtime=(1, 1))
+    assert_shape(res3[1].shape, IntTuple, runtime=(1, 3))
+    assert_shape(res3[2].shape, IntTuple, runtime=(2, 3))
 
 
 def test_atleast_3d() -> None:
@@ -242,15 +273,32 @@ def test_atleast_3d() -> None:
     assert_shape(jnp.atleast_3d(jnp.ones((2, 3, 4, 5))).shape, (2, 3, 4, 5))
     assert_shape(jnp.atleast_3d(np.ones((2, 3))).shape, (2, 3, 1))
 
-    # Zero arguments: returns list of arrays
+    # Zero arguments: returns empty tuple
     res0 = jnp.atleast_3d()
-    assert_type(res0, list[jax.Array])
-    assert res0 == []
+    assert_type(res0, tuple[()])
+    assert res0 == ()
 
-    # Two arguments: returns list of arrays
+    # Two arguments: returns 2-tuple of shape-typed arrays
     res2 = jnp.atleast_3d(jnp.ones(3), jnp.ones((2, 3)))
-    assert_type(res2, list[jax.Array])
-    assert isinstance(res2, list) and len(res2) == 2
+    assert isinstance(res2, tuple) and len(res2) == 2
     assert all(isinstance(x, jax.Array) for x in res2)
-    assert res2[0].shape == (1, 3, 1)
-    assert res2[1].shape == (2, 3, 1)
+    assert_shape(res2[0].shape, (1, 3, 1))
+    assert_shape(res2[1].shape, (2, 3, 1))
+
+    # Three array arguments: returns 3-tuple of shape-typed arrays
+    res3_arr = jnp.atleast_3d(jnp.ones(()), jnp.ones(3), jnp.ones((2, 3)))
+    assert_type(
+        res3_arr,
+        tuple[jax.Array[[1, 1, 1]], jax.Array[[1, 3, 1]], jax.Array[[2, 3, 1]]],
+    )
+    assert_shape(res3_arr[0].shape, (1, 1, 1))
+    assert_shape(res3_arr[1].shape, (1, 3, 1))
+    assert_shape(res3_arr[2].shape, (2, 3, 1))
+
+    # Three mixed arguments (with scalar): returns 3-tuple with gradual shapes
+    res3 = jnp.atleast_3d(1, jnp.ones(3), jnp.ones((2, 3)))
+    assert_type(res3, tuple[jax.Array, jax.Array, jax.Array])
+    assert isinstance(res3, tuple) and len(res3) == 3
+    assert_shape(res3[0].shape, IntTuple, runtime=(1, 1, 1))
+    assert_shape(res3[1].shape, IntTuple, runtime=(1, 3, 1))
+    assert_shape(res3[2].shape, IntTuple, runtime=(2, 3, 1))