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_


Discussion:
Thanks for the report, that's definitely a bug!

Until we get it fixed, you can just batch it with `jax.vmap`, as in `jax.vmap(jax.scipy.linalg.expm)(A)` (though also it doesn't accept int32, which seems to be another bug, so you'll need to explicitly cast like `jax.vmap(jax.scipy.linalg.expm)(A.astype('float32'))`).
I found via test failures that batched expm was only added in recent scipy versions, which explains why it wasn't part of the original JAX implementation 😁  (https://github.com/scipy/scipy/pull/15079)
I made some tests with `expm` and the result seems slower than I expect.

For a 2D array `x`,
```python
def expm(x):
    # option 1
    res = jnp.diag(jnp.exp(jnp.diag(x)))

    # option 2
    # TODO: vmap is very slow
    res = jax.vmap(jax.scipy.linalg.expm)(x[None, :, :])
    res = res[0]

    # option 3
    res = jax.scipy.linalg.expm(x)

    return res
```
This code was used in my project (Electromagnetic simulation) with jit option

option 1:
time at 1st run: 7.054819107055664
time at 2nd run: 0.008198022842407227

option 2:
time at 1st run: 106.56755089759827
time at 2nd run: 0.9272429943084717

option 3:
time at 1st run: 109.10750484466553
time at 2nd run: 0.758314847946167


I think it's slow even though I used `vmap` for a single operation.

For 3D array case,

```python
def expm_new(x):
    
    # option 1
    res = jnp.zeros(x.shape, dtype=jnp.complex128)
    ix = jnp.diag_indices_from(x[0])
    res = res.at[:, ix[0], ix[1]].set(jnp.exp(x[:, ix[0], ix[1]]))

    # option 2
    # res0 = jax.vmap(jax.scipy.linalg.expm)(x)   # TODO: remove vmap when the bug is fixed
    

    return res
```
option 1
time at 1st run:  0.32639575004577637
time at 2nd run: 0.00822305679321289

option 2
time at 1st run:  4.233128070831299
time at 2nd run: 0.92429518699646

Will it be improved or better to keep mine?

Thanks for the help!
If you're running microbenchmarks of JAX code, be sure to follow the advice at https://jax.readthedocs.io/en/latest/faq.html#benchmarking-jax-code (thinking about JIT compilation time, device transfer cost, asynchronous dispatch, etc.).

For example, running on a Colab CPU, I get these results (on GPU JAX is significantly faster, but it's hardly a fair comparison):
```python
import numpy as np
import scipy.linalg

import jax
import jax.numpy as jnp
from jax.scipy import linalg as jsp_linalg

# 2D Case
x_np = np.random.randn(200,200).astype('float32')
x_jnp = jnp.asarray(x_np)

jit_expm = jax.jit(jsp_linalg.expm)
_ = jit_expm(x_jnp)  # trigger compilation

%timeit scipy.linalg.expm(x_np)
# 15.3 ms ± 4.93 ms per loop (mean ± std. dev. of 7 runs, 100 loops each)
%timeit jit_expm(x_jnp).block_until_ready()
# 6.74 ms ± 1.05 ms per loop (mean ± std. dev. of 7 runs, 100 loops each)

# 3D Case
y_np = np.random.randn(3, 200,200).astype('float32')
y_jnp = jnp.asarray(y_np)

jit_vmap_expm = jax.jit(jax.vmap(jsp_linalg.expm))
_ = jit_vmap_expm(y_jnp)  # trigger compilation

%timeit scipy.linalg.expm(y_np)
# 56 ms ± 24.9 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)
%timeit jit_vmap_expm(y_jnp).block_until_ready()
# 53.8 ms ± 6.11 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)
```
Note that this is a situation (what amounts to basically a single LAPACK operation run on CPU) where I wouldn't necessarily expect JAX to be faster than SciPy; see [FAQ: is JAX faster than NumPy?](https://jax.readthedocs.io/en/latest/faq.html#is-jax-faster-than-numpy)
Thanks for the kind remind! I checked again but the same result. I was able to see the final calculation result in that time.

The process I used is bit more complex.

```python

@partial(jax.jit, static_argnums=(5, 6, 10, 11))
def field_dist_2d_vectorized_kji(wavelength, kx_vector, n_I, theta, phi, fourier_order_x, fourier_order_y, T1, layer_info_list, period,
                                 resolution=(10, 10, 10), type_complex=jnp.complex128):

    k0 = 2 * jnp.pi / wavelength

    fourier_indices_y = jnp.arange(-fourier_order_y, fourier_order_y + 1)
    ff_x = fourier_order_x * 2 + 1
    ff_y = fourier_order_y * 2 + 1
    ky_vector = k0 * (n_I * jnp.sin(theta) * jnp.sin(phi) + fourier_indices_y * (
            wavelength / period[1])).astype(type_complex)

    Kx = jnp.diag(jnp.tile(kx_vector, ff_y).flatten()) / k0
    Ky = jnp.diag(jnp.tile(ky_vector.reshape((-1, 1)), ff_x).flatten()) / k0

    resolution_z, resolution_y, resolution_x = resolution
    field_cell = jnp.zeros((resolution_z * len(layer_info_list), resolution_y, resolution_x, 6), dtype=type_complex)

    T_layer = T1

    big_I = jnp.eye((len(T1))).astype(type_complex)

    # From the first layer
    for idx_layer, (E_conv_i, q, W_11, W_12, W_21, W_22, V_11, V_12, V_21, V_22, big_X, big_A_i, big_B, d)\
            in enumerate(layer_info_list[::-1]):

        c = jnp.block([[big_I], [big_B @ big_A_i @ big_X]]) @ T_layer
        z_1d = jnp.arange(resolution_z).reshape((-1, 1, 1)) / resolution_z * d

        ff = len(c) // 4

        c1_plus = c[0 * ff:1 * ff]
        c2_plus = c[1 * ff:2 * ff]
        c1_minus = c[2 * ff:3 * ff]
        c2_minus = c[3 * ff:4 * ff]

        q1 = q[:len(q) // 2]
        q2 = q[len(q) // 2:]
        big_Q1 = jnp.diag(q1)
        big_Q2 = jnp.diag(q2)

        Sx = W_11 @ (expm_new(-k0 * big_Q1 * z_1d) @ c1_plus + expm_new(k0 * big_Q1 * (z_1d - d)) @ c1_minus) \
              + W_12 @ (expm_new(-k0 * big_Q2 * z_1d) @ c2_plus + expm_new(k0 * big_Q2 * (z_1d - d)) @ c2_minus)

        Sy = W_21 @ (expm_new(-k0 * big_Q1 * z_1d) @ c1_plus + expm_new(k0 * big_Q1 * (z_1d - d)) @ c1_minus) \
              + W_22 @ (expm_new(-k0 * big_Q2 * z_1d) @ c2_plus + expm_new(k0 * big_Q2 * (z_1d - d)) @ c2_minus)

        Ux = V_11 @ (-expm_new(-k0 * big_Q1 * z_1d) @ c1_plus + expm_new(k0 * big_Q1 * (z_1d - d)) @ c1_minus) \
              + V_12 @ (-expm_new(-k0 * big_Q2 * z_1d) @ c2_plus + expm_new(k0 * big_Q2 * (z_1d - d)) @ c2_minus)

        Uy = V_21 @ (-expm_new(-k0 * big_Q1 * z_1d) @ c1_plus + expm_new(k0 * big_Q1 * (z_1d - d)) @ c1_minus) \
              + V_22 @ (-expm_new(-k0 * big_Q2 * z_1d) @ c2_plus + expm_new(k0 * big_Q2 * (z_1d - d)) @ c2_minus)

        Sz = -1j * E_conv_i @ (Kx @ Uy - Ky @ Ux)
        Uz = -1j * (Kx @ Sy - Ky @ Sx)

        x_1d = jnp.arange(resolution_x).reshape((1, -1, 1))
        y_1d = jnp.arange(resolution_y).reshape((-1, 1, 1))

        x_1d = -1j * x_1d * period[0] / resolution_x
        y_1d = -1j * y_1d * period[1] / resolution_y

        x_2d = jnp.tile(x_1d, (resolution_y, 1, 1))
        y_2d = jnp.tile(y_1d, (1, resolution_x, 1))

        x_2d = x_2d * kx_vector
        y_2d = y_2d * ky_vector

        x_2d = x_2d.reshape((resolution_y, resolution_x, 1, len(kx_vector)))
        y_2d = y_2d.reshape((resolution_y, resolution_x, len(ky_vector), 1))

        exp_K = jnp.exp(x_2d) * jnp.exp(y_2d)
        exp_K = exp_K.reshape((resolution_y, resolution_x, -1))

        Ex = exp_K[:, :, None, :] @ Sx[:, None, None, :, :]
        Ey = exp_K[:, :, None, :] @ Sy[:, None, None, :, :]
        Ez = exp_K[:, :, None, :] @ Sz[:, None, None, :, :]

        Hx = -1j * exp_K[:, :, None, :] @ Ux[:, None, None, :, :]
        Hy = -1j * exp_K[:, :, None, :] @ Uy[:, None, None, :, :]
        Hz = -1j * exp_K[:, :, None, :] @ Uz[:, None, None, :, :]

        val = jnp.concatenate(
            (Ex.squeeze(-1), Ey.squeeze(-1), Ez.squeeze(-1), Hx.squeeze(-1), Hy.squeeze(-1), Hz.squeeze(-1)),
            axis=-1)

        field_cell = field_cell.at[resolution_z * idx_layer:resolution_z * (idx_layer + 1)].set(val)

        T_layer = big_A_i @ big_X @ T_layer

    return field_cell

```

There is extensive matrix operation so I think this can be a reason that differs from your test result. Could be a compilation optimization thing?

Full code is uploaded here
https://github.com/kc-ml2/meent/blob/DEV_JAX_expm/benchmarks/jax_expm.py

and expm option is here
https://github.com/kc-ml2/meent/blob/DEV_JAX_expm/meent/on_jax/emsolver/field_distribution.py

```python
def expm(x):
    # option 1
    res = jnp.diag(jnp.exp(jnp.diag(x)))

    # option 2
    # res = jax.vmap(jax.scipy.linalg.expm)(x[None, :, :]) # vmap is very slow
    # res = res[0]

    # option 3
    # res = jax.scipy.linalg.expm(x)

    return res


def expm_new(x):
    # option 1
    res = jnp.zeros(x.shape, dtype=jnp.complex128)
    ix = jnp.diag_indices_from(x[0])
    res = res.at[:, ix[0], ix[1]].set(jnp.exp(x[:, ix[0], ix[1]]))
    # option 2
    # res = jax.vmap(jax.scipy.linalg.expm)(x)  # remove vmap when the bug is fixed. But vmap is very slow.

    return res
```

Thanks for the full example, but I'm unsure what your question is.
> Note that this is a situation (what amounts to basically a single LAPACK operation run on CPU) where I wouldn't necessarily expect JAX to be faster than SciPy; see [FAQ: is JAX faster than NumPy?](https://jax.readthedocs.io/en/latest/faq.html#is-jax-faster-than-numpy)

I just want to stress that option 1 (manual JAX expm) is much faster than option 2 (SciPy expm).

```python
def expm_new(x):
    
    # option 1
    res = jnp.zeros(x.shape, dtype=jnp.complex128)
    ix = jnp.diag_indices_from(x[0])
    res = res.at[:, ix[0], ix[1]].set(jnp.exp(x[:, ix[0], ix[1]]))

    # option 2
    # res0 = jax.vmap(jax.scipy.linalg.expm)(x)   # TODO: remove vmap when the bug is fixed
    
    return res
```
Thanks - it makes sense that in a jit-compiled sequence of operations in JAX would be faster than a similar sequence of operations in numpy/scipy.

Also, if you're directly comparing jax performance to scipy performance, differing dtypes might play a role. numpy/scipy default to float64 operations, while JAX will do everything in float32 unless you explicitly enable float64 (see [Sharp Bits: 64-bit Precision](https://jax.readthedocs.io/en/latest/notebooks/Common_Gotchas_in_JAX.html#double-64bit-precision)).
yep. That was the same condition - 64bit.
Thanks!
Please be careful of the fact that the `jax_enable_x64` config should be set at the start of the script, before running any other JAX code. This is discussed at the link I shared above.
OK. Thanks for the kind remind :) 
But now I'm confused.
>Thanks - it makes sense that in a jit-compiled sequence of operations in JAX would be faster than a similar sequence of operations in numpy/scipy.  

Then is it recommended to write manual JAX function rather using existing jax.numpy or jax.scipy in case of jit-compiled sequence?
Sorry, I think I'm confused about what you're asking. Which two operations are you comparing? I thought you were comparing `scipy.linalg.expm` to `jax.scipy.linalg.expm`. Is that not the case?
Yes I thought there was miscommunication :)

I compared JAX manual expm function and `jax.scipy.linalg.expm`.

```python
    # option 1
    res = jnp.zeros(x.shape, dtype=jnp.complex128)
    ix = jnp.diag_indices_from(x[0])
    res = res.at[:, ix[0], ix[1]].set(jnp.exp(x[:, ix[0], ix[1]]))

    # option 2
    res0 = jax.vmap(jax.scipy.linalg.expm)(x)   # TODO: remove vmap when the bug is fixed
```
option 1 and 2 are `expm` function but option 1 is what I wrote and option 2 is from jax.scipy.


> option 1
> time at 1st run: 0.32639575004577637
> time at 2nd run: 0.00822305679321289
> 
> option 2
> time at 1st run: 4.233128070831299
> time at 2nd run: 0.92429518699646

and option 1 is much faster.
I see, thanks for the clarification. Those are entirely different operations, The first is extracting the diagonal component of the matrix and computing its exponent. This is a relatively fast computation. The second computes the full expm of the matrix. This is a relatively costly computation.

These are entirely different operations, but they will agree if `x` happens to be a diagonal matrix. If you know *a priori* that the input is a diagonal matrix, you should use the faster approach. If you do not know this *a priori* and your intent is to compute the matrix exponent, then the faster approach will in general return the incorrect result.
Oh my... You are right. I misunderstood `expm`.
Really appreciate for your kindness.

Resolving patch:
diff --git a/jax/_src/scipy/linalg.py b/jax/_src/scipy/linalg.py
--- a/jax/_src/scipy/linalg.py
+++ b/jax/_src/scipy/linalg.py
@@ -433,6 +433,16 @@ def triu(m: ArrayLike, k: int = 0) -> Array:
 @_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):


Test patch:
diff --git a/tests/linalg_test.py b/tests/linalg_test.py
--- a/tests/linalg_test.py
+++ b/tests/linalg_test.py
@@ -29,12 +29,15 @@
 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 @@ def testTriangularSolveGradPrecision(self):
 
   @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)