Fix jax stub compatibility

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

@@ -50,7 +50,7 @@ from shape_extensions import (
 type _Shape = IntTuple
 type _Axis = int | tuple[int, ...] | None
 type _SingleAxis = int | None
-type _Scalar = bool | int | float | complex | np.number
+type _Scalar = bool | int | float | complex | np.number | np.bool_
 # Note: when using _ArrayLike in an annotation, the Shape passed to it must
 # have a default value of [] for scalars to be handled properly.
 type ArrayLike[Shape: _Shape] = Array[Shape] | np.ndarray[Shape] | _Scalar

tensor-shapes/jax-pyrefly-stubs/jax-stubs/_src/basearray.pyi

@@ -0,0 +1,10 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+import numpy as np
+from jax._array import Array as Array
+
+type StaticScalar = np.bool_ | np.number | bool | int | float | complex
+type ArrayLike = Array | np.ndarray | StaticScalar

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

@@ -42,6 +42,15 @@ def test_elementwise_operators_preserve_shape() -> None:
     assert_shape((b**2).shape, (3, 4))
 
 
+def test_numpy_boolean_scalar() -> None:
+    a = jnp.ones((2, 3))
+
+    assert_shape((a * np.True_).shape, (2, 3))
+    assert_shape((np.True_ * a).shape, (2, 3))
+    assert_shape(jnp.multiply(a, np.False_).shape, (2, 3))
+    np.testing.assert_allclose(np.True_ * a, np.ones((2, 3)))
+
+
 def test_scalar_operands_preserve_shape() -> None:
     a = jnp.ones((3, 4))
     v = jnp.full(5, 2.0)

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

@@ -7,6 +7,7 @@
 
 import jax
 import jax.numpy as jnp
+import jax.scipy.linalg as jsp_linalg
 import numpy as np
 from shape_extensions import assert_shape, IntTuple, IntVar
 
@@ -117,6 +118,23 @@ def test_solve_matrix_rhs() -> None:
     assert_shape(jnp.linalg.solve(a, b).shape, (3, 2))
 
 
+def test_native_triangular_solve() -> None:
+    a = jnp.eye(3)
+    b = jnp.ones(3)
+    x: jax.Array = jsp_linalg.solve_triangular(a, b)
+
+    assert_shape(x.shape, IntTuple, runtime=(3,))
+    np.testing.assert_allclose(a @ x, b)
+
+    try:
+        # E: is not assignable to parameter `a`
+        jsp_linalg.solve_triangular(object(), b)
+    except TypeError:
+        pass
+    else:
+        raise AssertionError("expected JAX to reject a non-array coefficient matrix")
+
+
 def test_inv_and_matrix_power() -> None:
     a = jnp.eye(4)