Repository: google/jax
Issue: expm does not support ndarray

### Description

Hi, I tried to use jax.scipy.linalg.expm with 3 dimensional array but got error.
Docs says it's doable when last 2 dimensions have same length but failed. Can you please check?

https://jax.readthedocs.io/en/latest/_autosummary/jax.scipy.linalg.expm.html

```python
import jax
import jax.numpy as jnp
A = jnp.arange(2*11*11).reshape((2,11,11))
jax.scipy.linalg.expm(A)
```

> Traceback (most recent call last):
>   File "/Users/yongha/miniforge3_m1/envs/rcwa/lib/python3.9/site-packages/IPython/core/interactiveshell.py", line 3378, in run_code
>     exec(code_obj, self.user_global_ns, self.user_ns)
>   File "<ipython-input-5-d2ac0f4d7ead>", line 1, in <module>
>     jax.scipy.linalg.expm(A)
>   File "/Users/yongha/miniforge3_m1/envs/rcwa/lib/python3.9/site-packages/jax/_src/traceback_util.py", line 162, in reraise_with_filtered_traceback
>     return fun(*args, **kwargs)
>   File "/Users/yongha/miniforge3_m1/envs/rcwa/lib/python3.9/site-packages/jax/_src/api.py", line 622, in cache_miss
>     execute = dispatch._xla_call_impl_lazy(fun_, *tracers, **params)
>   File "/Users/yongha/miniforge3_m1/envs/rcwa/lib/python3.9/site-packages/jax/_src/dispatch.py", line 241, in _xla_call_impl_lazy
>     return xla_callable(fun, device, backend, name, donated_invars, keep_unused,
>   File "/Users/yongha/miniforge3_m1/envs/rcwa/lib/python3.9/site-packages/jax/linear_util.py", line 303, in memoized_fun
>     ans = call(fun, *args)
>   File "/Users/yongha/miniforge3_m1/envs/rcwa/lib/python3.9/site-packages/jax/_src/dispatch.py", line 357, in _xla_callable_uncached
>     computation = sharded_lowering(fun, device, backend, name, donated_invars,
>   File "/Users/yongha/miniforge3_m1/envs/rcwa/lib/python3.9/site-packages/jax/_src/dispatch.py", line 348, in sharded_lowering
>     return pxla.lower_sharding_computation(
>   File "/Users/yongha/miniforge3_m1/envs/rcwa/lib/python3.9/site-packages/jax/_src/profiler.py", line 314, in wrapper
>     return func(*args, **kwargs)
>   File "/Users/yongha/miniforge3_m1/envs/rcwa/lib/python3.9/site-packages/jax/interpreters/pxla.py", line 2792, in lower_sharding_computation
>     jaxpr, global_out_avals, consts = pe.trace_to_jaxpr_final(
>   File "/Users/yongha/miniforge3_m1/envs/rcwa/lib/python3.9/site-packages/jax/_src/profiler.py", line 314, in wrapper
>     return func(*args, **kwargs)
>   File "/Users/yongha/miniforge3_m1/envs/rcwa/lib/python3.9/site-packages/jax/interpreters/partial_eval.py", line 2065, in trace_to_jaxpr_final
>     jaxpr, out_avals, consts = trace_to_subjaxpr_dynamic(
>   File "/Users/yongha/miniforge3_m1/envs/rcwa/lib/python3.9/site-packages/jax/interpreters/partial_eval.py", line 1998, in trace_to_subjaxpr_dynamic
>     ans = fun.call_wrapped(*in_tracers_)
>   File "/Users/yongha/miniforge3_m1/envs/rcwa/lib/python3.9/site-packages/jax/linear_util.py", line 167, in call_wrapped
>     ans = self.f(*args, **dict(self.params, **kwargs))
>   File "/Users/yongha/miniforge3_m1/envs/rcwa/lib/python3.9/site-packages/jax/_src/scipy/linalg.py", line 436, in expm
>     P, Q, n_squarings = _calc_P_Q(A)
>   File "/Users/yongha/miniforge3_m1/envs/rcwa/lib/python3.9/site-packages/jax/_src/traceback_util.py", line 162, in reraise_with_filtered_traceback
>     return fun(*args, **kwargs)
>   File "/Users/yongha/miniforge3_m1/envs/rcwa/lib/python3.9/site-packages/jax/_src/api.py", line 626, in cache_miss
>     top_trace.process_call(primitive, fun_, tracers, params))
>   File "/Users/yongha/miniforge3_m1/envs/rcwa/lib/python3.9/site-packages/jax/interpreters/partial_eval.py", line 1739, in process_call
>     jaxpr, out_type, consts = trace_to_subjaxpr_dynamic2(f, self.main, debug_info=dbg)
>   File "/Users/yongha/miniforge3_m1/envs/rcwa/lib/python3.9/site-packages/jax/interpreters/partial_eval.py", line 2027, in trace_to_subjaxpr_dynamic2
>     ans = fun.call_wrapped(*in_tracers_)
>   File "/Users/yongha/miniforge3_m1/envs/rcwa/lib/python3.9/site-packages/jax/linear_util.py", line 167, in call_wrapped
>     ans = self.f(*args, **dict(self.params, **kwargs))
>   File "/Users/yongha/miniforge3_m1/envs/rcwa/lib/python3.9/site-packages/jax/_src/scipy/linalg.py", line 455, in _calc_P_Q
>     raise ValueError('expected A to be a square matrix')
> ValueError: expected A to be a square matrix

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

jax 0.4.1,  jaxlib 0.4.1

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

CPU

### Additional system info

M1

### NVIDIA GPU info

_No response_

Changed files:
tests/linalg_test.py, jax/_src/scipy/linalg.py

diff --git a/jax/_src/scipy/linalg.py b/jax/_src/scipy/linalg.py
index 4b6026a23..91db1885e 100644
--- a/jax/_src/scipy/linalg.py
+++ b/jax/_src/scipy/linalg.py
@@ -433,6 +433,16 @@ where norm() denotes the L1 norm, and
 @_wraps(scipy.linalg.expm, lax_description=_expm_description)
 @partial(jit, static_argnames=('upper_triangular', 'max_squarings'))
 def expm(A: ArrayLike, *, upper_triangular: bool = False, max_squarings: int = 16) -> Array:
+  A, = promote_dtypes_inexact(A)
+
+  if A.ndim < 2 or A.shape[-1] != A.shape[-2]:
+    raise ValueError(f"Expected A to be a (batched) square matrix, got {A.shape=}.")
+
+  if A.ndim > 2:
+    return jnp.vectorize(
+      partial(expm, upper_triangular=upper_triangular, max_squarings=max_squarings),
+      signature="(n,n)->(n,n)")(A)
+
   P, Q, n_squarings = _calc_P_Q(A)
 
   def _nan(args):
diff --git a/tests/linalg_test.py b/tests/linalg_test.py
index 7534118fe..b5c950135 100644
--- a/tests/linalg_test.py
+++ b/tests/linalg_test.py
@@ -29,12 +29,15 @@ from jax import jit, grad, jvp, vmap
 from jax import lax
 from jax import numpy as jnp
 from jax import scipy as jsp
+from jax._src.numpy.util import promote_dtypes_inexact
 from jax._src import test_util as jtu
 
 from jax.config import config
 config.parse_flags_with_absl()
 FLAGS = config.FLAGS
 
+scipy_version = tuple(map(int, scipy.version.version.split('.')[:3]))
+
 T = lambda x: np.swapaxes(x, -1, -2)
 
 
@@ -1240,18 +1243,21 @@ class ScipyLinalgTest(jtu.JaxTestCase):
 
   @jtu.sample_product(
     n=[1, 4, 5, 20, 50, 100],
-    dtype=float_types + complex_types,
+    batch_size=[(), (2,), (3, 4)] if scipy_version >= (1, 9, 0) else [()],
+    dtype=int_types + float_types + complex_types
   )
-  def testExpm(self, n, dtype):
+  def testExpm(self, n, batch_size, dtype):
     rng = jtu.rand_small(self.rng())
-    args_maker = lambda: [rng((n, n), dtype)]
+    args_maker = lambda: [rng((*batch_size, n, n), dtype)]
 
-    osp_fun = lambda a: osp.linalg.expm(a)
-    jsp_fun = lambda a: jsp.linalg.expm(a)
+    # Compare to numpy with JAX type promotion semantics.
+    def osp_fun(A):
+      return osp.linalg.expm(np.array(*promote_dtypes_inexact(A)))
+    jsp_fun = jsp.linalg.expm
     self._CheckAgainstNumpy(osp_fun, jsp_fun, args_maker)
     self._CompileAndCheck(jsp_fun, args_maker)
 
-    args_maker_triu = lambda: [np.triu(rng((n, n), dtype))]
+    args_maker_triu = lambda: [np.triu(rng((*batch_size, n, n), dtype))]
     jsp_fun_triu = lambda a: jsp.linalg.expm(a, upper_triangular=True)
     self._CheckAgainstNumpy(osp_fun, jsp_fun_triu, args_maker_triu)
     self._CompileAndCheck(jsp_fun_triu, args_maker_triu)