BUG: jnp.linalg.matrix_rank does not operate on stacks of matrices
### Description

```
myInput = np.random.normal(size=(9,10,10))
myInput.shape
# (9,10,10)
np.linalg.matrix_rank(myInput)
# array([10, 10, 10, 10, 10, 10, 10, 10, 10])
```

jax fails on this same stack:
```
jax.numpy.linalg.matrix_rank(jnp.array(myInput))
### TypeError: array should have 2 or fewer dimensions
### Expected: DeviceArray([10, 10, 10, 10, 10, 10, 10, 10, 10])
```

### What jax/jaxlib version are you using?

jax 0.3.13, jaxlib 0.3.10

### Which accelerator(s) are you using?

- [ ] CPU
- [X] GPU
- [ ] TPU

### Additional System Info

_No response_
diff --git a/jax/_src/numpy/linalg.py b/jax/_src/numpy/linalg.py
index 3f3b172bc..97827576e 100644
--- a/jax/_src/numpy/linalg.py
+++ b/jax/_src/numpy/linalg.py
@@ -111,14 +111,13 @@ def matrix_power(a, n):
 @jit
 def matrix_rank(M, tol=None):
   M, = _promote_dtypes_inexact(jnp.asarray(M))
-  if M.ndim > 2:
-    raise TypeError("array should have 2 or fewer dimensions")
   if M.ndim < 2:
     return jnp.any(M != 0).astype(jnp.int32)
   S = svd(M, full_matrices=False, compute_uv=False)
   if tol is None:
-    tol = S.max() * np.max(M.shape).astype(S.dtype) * jnp.finfo(S.dtype).eps
-  return jnp.sum(S > tol)
+    tol = S.max(-1) * np.max(M.shape[-2:]).astype(S.dtype) * jnp.finfo(S.dtype).eps
+  tol = jnp.expand_dims(tol, np.ndim(tol))
+  return jnp.sum(S > tol, axis=-1)
 
 
 @custom_jvp
diff --git a/tests/linalg_test.py b/tests/linalg_test.py
index 97234906f..ff9f28c95 100644
--- a/tests/linalg_test.py
+++ b/tests/linalg_test.py
@@ -924,7 +924,8 @@ class NumpyLinalgTest(jtu.JaxTestCase):
       {"testcase_name": "_shape={}".format(
            jtu.format_shape_dtype_string(shape, dtype)),
        "shape": shape, "dtype": dtype}
-      for shape in [(3, ), (1, 2), (8, 5), (4, 4), (5, 5), (50, 50)]
+      for shape in [(3, ), (1, 2), (8, 5), (4, 4), (5, 5), (50, 50),
+                    (3, 4, 5), (2, 3, 4, 5)]
       for dtype in float_types + complex_types))
   def testMatrixRank(self, shape, dtype):
     rng = jtu.rand_default(self.rng())
['tests/linalg_test.py', 'jax/_src/numpy/linalg.py']