Repository: jax-ml/jax (python)

Issue:
`jax.numpy.linalg.matrix_power` brings different results with `numpy.linalg.matrix_power`
### Description

I found that `jax.numpy.linalg.matrix_power` brings different results with `numpy.linalg.matrix_power` in following case.
I think this may be a bug in `JAX`.

```py
import numpy as np
import jax.numpy as jnp

print(np.linalg.matrix_power(np.array([[True,True], [False,True]]), 2))
print(jnp.linalg.matrix_power(jnp.array([[True, True], [False, True]]), 2))
```

Output:
```bash
[[ True  True]
 [False  True]]
[[1. 2.]
 [0. 1.]]
```

### System info (python version, jaxlib version, accelerator, etc.)

```
jax:    0.6.0
jaxlib: 0.6.0
numpy:  2.2.3
python: 3.10.12 (main, Feb  4 2025, 14:57:36) [GCC 11.4.0]
device info: NVIDIA GeForce RTX 4090-1, 1 local devices"
process_count: 1
platform: uname_result(system='Linux', node='4db45dc420f8', release='6.11.0-25-generic', version='#25~24.04.1-Ubuntu SMP PREEMPT_DYNAMIC Tue Apr 15 17:20:50 UTC 2', machine='x86_64')


$ nvidia-smi
Wed May  7 16:26:08 2025       
+---------------------------------------------------------------------------------------+
| NVIDIA-SMI 535.230.02             Driver Version: 535.230.02   CUDA Version: 12.2     |
|-----------------------------------------+----------------------+----------------------+
| GPU  Name                 Persistence-M | Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp   Perf          Pwr:Usage/Cap |         Memory-Usage | GPU-Util  Compute M. |
|                                         |                      |               MIG M. |
|=========================================+======================+======================|
|   0  NVIDIA GeForce RTX 4090        Off | 00000000:01:00.0 Off |                  Off |
| 36%   32C    P2              36W / 450W |    407MiB / 24564MiB |      1%      Default |
|                                         |                      |                  N/A |
+-----------------------------------------+----------------------+----------------------+
                                                                                         
+---------------------------------------------------------------------------------------+
| Processes:                                                                            |
|  GPU   GI   CI        PID   Type   Process name                            GPU Memory |
|        ID   ID                                                             Usage      |
|=======================================================================================|
+---------------------------------------------------------------------------------------+
```

Pull request description:
jnp.linalg.matrix_power: support non-float inputs
Fixes #28603

Patch:
diff --git a/jax/_src/numpy/linalg.py b/jax/_src/numpy/linalg.py
index 146bbbda0..0e20e5b2a 100644
--- a/jax/_src/numpy/linalg.py
+++ b/jax/_src/numpy/linalg.py
@@ -367,8 +367,7 @@ def matrix_power(a: ArrayLike, n: int) -> Array:
     Array([[ 5.5 , -2.5 ],
            [-3.75,  1.75]], dtype=float32)
   """
-  a = ensure_arraylike("jnp.linalg.matrix_power", a)
-  arr, = promote_dtypes_inexact(a)
+  arr = ensure_arraylike("jnp.linalg.matrix_power", a)
 
   if arr.ndim < 2:
     raise TypeError("{}-dimensional array given. Array must be at least "


Test patch:
diff --git a/tests/linalg_test.py b/tests/linalg_test.py
index 74259c300..033ca989c 100644
--- a/tests/linalg_test.py
+++ b/tests/linalg_test.py
@@ -1230,6 +1230,13 @@ class NumpyLinalgTest(jtu.JaxTestCase):
     self._CompileAndCheck(partial(jnp.linalg.matrix_power, n=n), args_maker,
                           rtol=1e-3)
 
+  def testMatrixPowerBool(self):
+    # Regression test for https://github.com/jax-ml/jax/issues/28603
+    mat = np.array([[True,True], [False,True]])
+    np_result = np.linalg.matrix_power(mat, 2)
+    jnp_result = jnp.linalg.matrix_power(mat, 2)
+    self.assertArraysEqual(np_result, jnp_result)
+
   @jtu.sample_product(
     shape=[(3, ), (1, 2), (8, 5), (4, 4), (5, 5), (50, 50), (3, 4, 5),
            (2, 3, 4, 5)],