diff --git a/tensor-shapes/jax-pyrefly-stubs/jax-stubs/__init__.pyi b/tensor-shapes/jax-pyrefly-stubs/jax-stubs/__init__.pyi index d51cc49f28..0849bdfc01 100644 --- a/tensor-shapes/jax-pyrefly-stubs/jax-stubs/__init__.pyi +++ b/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 ( diff --git a/tensor-shapes/jax-pyrefly-stubs/jax-stubs/_shapes.pyi b/tensor-shapes/jax-pyrefly-stubs/jax-stubs/_shapes.pyi index 03b5be68ea..f7a08f8017 100644 --- a/tensor-shapes/jax-pyrefly-stubs/jax-stubs/_shapes.pyi +++ b/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( diff --git a/tensor-shapes/jax-pyrefly-stubs/jax-stubs/numpy/__init__.pyi b/tensor-shapes/jax-pyrefly-stubs/jax-stubs/numpy/__init__.pyi index 6a17a29146..195f905538 100644 --- a/tensor-shapes/jax-pyrefly-stubs/jax-stubs/numpy/__init__.pyi +++ b/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 diff --git a/tensor-shapes/jax-pyrefly-stubs/jax-stubs/random.pyi b/tensor-shapes/jax-pyrefly-stubs/jax-stubs/random.pyi index 70ca749b03..886073083e 100644 --- a/tensor-shapes/jax-pyrefly-stubs/jax-stubs/random.pyi +++ b/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], diff --git a/tensor-shapes/jax-pyrefly-stubs/test/test_random.py b/tensor-shapes/jax-pyrefly-stubs/test/test_random.py index b50a62b286..0f41e79e08 100644 --- a/tensor-shapes/jax-pyrefly-stubs/test/test_random.py +++ b/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) diff --git a/tensor-shapes/jax-pyrefly-stubs/test/test_reshape.py b/tensor-shapes/jax-pyrefly-stubs/test/test_reshape.py index 62d9832393..e345fee2a6 100644 --- a/tensor-shapes/jax-pyrefly-stubs/test/test_reshape.py +++ b/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)) diff --git a/tensor-shapes/jax-pyrefly-stubs/test/test_typing.py b/tensor-shapes/jax-pyrefly-stubs/test/test_typing.py index 966a228bc3..79f6dbb39d 100644 --- a/tensor-shapes/jax-pyrefly-stubs/test/test_typing.py +++ b/tensor-shapes/jax-pyrefly-stubs/test/test_typing.py @@ -45,6 +45,29 @@ def test_dtype_like() -> None: assert_shape(jnp.ones((2,), dtype=jnp.float32).shape, (2,)) +def test_scalar_type_protocol() -> None: + assert_type(jnp.float32.dtype, np.dtype[Any]) + assert_type(jnp.float32.__numpy_dtype__, np.dtype[Any]) + assert jnp.float32.dtype == np.dtype("float32") + assert jnp.float32.__numpy_dtype__ == np.dtype("float32") + + x = jnp.float32(1.5) + assert_type(x, jax.Array) + assert_shape(x.shape, IntTuple, runtime=()) + + y = jnp.int32([1, 2, 3]) + assert_type(y, jax.Array) + assert_shape(y.shape, IntTuple, runtime=(3,)) + + assert_type(hash(jnp.float32), int) + assert_type(jnp.float32 == jnp.float32, bool) + assert jnp.float32 == jnp.float32 + assert jnp.float32 != jnp.int32 + + with assert_raises(TypeError): + jnp.float32() # E: Missing positional argument + + def test_scalar_conversions() -> None: scalar = jnp.array(3.0) assert_shape(scalar.shape, ()) diff --git a/tensor-shapes/test-requirements.txt b/tensor-shapes/test-requirements.txt index 00a56a21e5..571999b1e9 100644 --- a/tensor-shapes/test-requirements.txt +++ b/tensor-shapes/test-requirements.txt @@ -19,10 +19,10 @@ numpy==2.5.0 einops==0.8.2 -jax==0.11.2 +jax==0.12.0 # Pinned alongside `jax`: it is a separate distribution that would otherwise # resolve freely under a pinned `jax`. -jaxlib==0.11.2 +jaxlib==0.12.0 torch==2.12.1+cpu ; sys_platform != "darwin" torch==2.12.1 ; sys_platform == "darwin"