From d949e3b047b7d6004a1ff501215f1cdb1712bd2d Mon Sep 17 00:00:00 2001
From: Brodi Elwood <brodielwood@gmail.com>
Date: Tue, 20 Oct 2026 21:08:17 -0400
Subject: [PATCH] Fix jax array compatibility

---
 .../jax-pyrefly-stubs/jax-stubs/_array.pyi     |  2 +-
 .../jax-stubs/_src/basearray.pyi               | 10 ++++++++++
 .../jax-pyrefly-stubs/test/test_arithmetic.py  |  9 +++++++++
 .../jax-pyrefly-stubs/test/test_linalg.py      | 18 ++++++++++++++++++
 4 files changed, 38 insertions(+), 1 deletion(-)
 create mode 100644 tensor-shapes/jax-pyrefly-stubs/jax-stubs/_src/basearray.pyi

diff --git a/tensor-shapes/jax-pyrefly-stubs/jax-stubs/_array.pyi b/tensor-shapes/jax-pyrefly-stubs/jax-stubs/_array.pyi
index 954aa73cb1..34aa12e963 100644
--- a/tensor-shapes/jax-pyrefly-stubs/jax-stubs/_array.pyi
+++ b/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
diff --git a/tensor-shapes/jax-pyrefly-stubs/jax-stubs/_src/basearray.pyi b/tensor-shapes/jax-pyrefly-stubs/jax-stubs/_src/basearray.pyi
new file mode 100644
index 0000000000..be60a629bf
--- /dev/null
+++ b/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
diff --git a/tensor-shapes/jax-pyrefly-stubs/test/test_arithmetic.py b/tensor-shapes/jax-pyrefly-stubs/test/test_arithmetic.py
index 0f11c85063..777371d77b 100644
--- a/tensor-shapes/jax-pyrefly-stubs/test/test_arithmetic.py
+++ b/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)
diff --git a/tensor-shapes/jax-pyrefly-stubs/test/test_linalg.py b/tensor-shapes/jax-pyrefly-stubs/test/test_linalg.py
index 1dd403ae49..765a42ba7f 100644
--- a/tensor-shapes/jax-pyrefly-stubs/test/test_linalg.py
+++ b/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)