Repository: scipy/scipy

Issue:
Docs suggest scipy.sparse.linalg.expm_multiply supports LinearOperator but it doesn't
According to [the docs](http://docs.scipy.org/doc/scipy/reference/generated/scipy.sparse.linalg.expm_multiply.html#scipy.sparse.linalg.expm_multiply) the first parameter to expm_multiply is a "transposable linear operator", which presumably includes LinearOperator. However expm_multiply actually appears to accept numpy arrays or scipy sparse matrices only. For example:

``` python
import numpy as np
import scipy.sparse.linalg as spsl
spsl.expm_multiply(spsl.aslinearoperator(np.array([[2]])), np.array([1]))
```

fails because [_trace()](https://github.com/scipy/scipy/blob/v0.16.0/scipy/sparse/linalg/_expm_multiply.py#L31) assumes its parameter is either a numpy array or a scipy sparse matrix.

It seems expm_multiply uses the trace, 1-norm and infinity-norm of its parameter, none of which are available for LinearOperator. <del>Fixing this bug by supporting LinearOperator would therefore require some non-trivial algorithm design so you may want to fix the documentation now and defer adding LinearOperator support until later.</del>



Discussion:
Actually this bug is easier to fix that it appeared at first. Let `M` be the matrix being exponentiated. There’s already code for estimating the 1-norm using only matrix-vector products (which is applied to powers of `M`) so unless exact one-norm is needed for `M` itself that should work fine. The infinity-norm isn’t used on `M` so that’s not an issue. The biggest issue appears to be computing the trace of `M`. I can think of two ways to find the trace. One option is to ask `LinearOperator` implementations used with `expm_multiply` to provide a `trace()` method. I’d guess that many `LinearOperator` implementations can compute the trace about as easily as they can compute matrix/vector products; for example it’s easy to find the trace of matrices with a known low rank representation. If an approximate trace is ok (which seems likely but I haven’t read the papers) another option is to use the following unbiased estimator of the trace:

``` python
import math
import numpy as np
def EstimatedTrace(A, k):
  n = A.shape[1]
  rand_mat = np.random.randn(n, k)
  # Normalize columns of rand_mat to have norm sqrt(n)
  rand_mat *= (math.sqrt(n) / np.linalg.norm(rand_mat, axis=0))
  # TODO: if A.dot(rand_mat) returns a np.matrix or scipy.sparse matrix
  # the multiplication won't be compentwise as desired
  return (A.dot(rand_mat) * rand_mat).sum()/k
```

The trace can also be computed using (naive implementation)

``` py
def trace(x):
    tr = 0
    for i in xrange(x.shape[0]):
        v = np.zeros(x.shape[0])
        v[i] = 1
        tr += x.dot(v)[i]
    return tr
```

I'm not sure if this would be an acceptable fallback.

Ping @argriffing.

@wschudy I see you're right, it doesn't work for abstract linear operators so the docs should be changed.  Looking at the publication http://eprints.ma.man.ac.uk/1591/01/covered/MIMS_ep2010_30.pdf I see that the algorithm does ask for the trace.  This is efficient for already-dense matrices, and scipy.sparse matrices have `.diagonal()`.  I understand and agree with the point that for a low rank linear operator the trace can be relatively efficiently computed (because the trace is the sum of eigenvalues, and a low rank matrix has only a few nonzero eigenvalues).  I don't think scipy has an interface for abstract linear operators with a known trace or spectrum; for example https://github.com/scipy/scipy/blob/master/scipy/sparse/linalg/interface.py doesn't include 'trace' or 'eigen'-anything.

@larsmans To be honest, I'm not caught up with the latest scipy linear operator 'best practices'.  For example I wasn't completely following the details of https://github.com/scipy/scipy/pull/4087, and I didn't even try to follow the relevant numpy mega-thread https://github.com/numpy/numpy/issues/5844. In my opinion it would be nice if scipy would adopt some battle-tested abstract linear operator interface, but of course this may require more volunteer work than is available.

1. Looking at the paper, the trace is only used as a heuristic to find a shift mu such that || (A - mu*I)^p ||_1 is small, which reduces runtime. So an approximate trace (or an alternate heuristic to find mu) is fine.
2. @larsmans: using n matrix-vector multiplies to calculate the trace would dominate the overall runtime for many applications. An approximation is better.
3. IIRC the limit as p goes to infinity of || M^p ||^(1/p) is the spectral_radius(M) for all matrix norms. So one reasonable heuristic for mu would be the average of the smallest and largest eigenvalue as estimated using Arnoldi iteration with big tolerances. The existing scipy.sparse.linalg.eigs doesn't have the right interface for this but for this application most of the fanciness in ARPACK may not be needed and naive Arnoldi may work? For symmetric matrices at least it looks like O~(sqrt ln n) Arnoldi (= Lanczos) iterations suffices to get the spectral radius of A - mu \* I  to within a constant factor of optimal. Proof sketch:  O~(sqrt ln n) Lanczos iterations is enough for the Krylov subspace to approximate exp(M ln n) [well](http://www.sciencedirect.com/science/article/pii/0377042796000064), where M is A shifted and scaled to have min eigenvalue -1 and max eigenvalue 1. That implies the Krylov subspace also approximates the max eigenvalue well.

I've recently run into this problem, and I'm wondering if any progress has been made since this issue was raised. It seems like giving LinearOperator the option to define it's own trace should solve the problem (but I'm not very experienced with this interface). I'm using Scipy version 1.1.0.
Hello from the future! I have also just run into this problem. I'm using Scipy version 1.5.1.  Are there any simple workarounds?
I run into this problem also. expm_multply should be useful in many occasions, provided an abstract LinearOperator interlace is supported.   @argriffing 
This seems to still be an issue nearly 6 years later - are there any plans to fix this in the next version? @argriffing 

Resolving patch:
diff --git a/scipy/sparse/linalg/_expm_multiply.py b/scipy/sparse/linalg/_expm_multiply.py
--- a/scipy/sparse/linalg/_expm_multiply.py
+++ b/scipy/sparse/linalg/_expm_multiply.py
@@ -1,11 +1,15 @@
-"""Compute the action of the matrix exponential.
-"""
+"""Compute the action of the matrix exponential."""
+from warnings import warn
 
 import numpy as np
 
-import scipy.sparse
-from scipy.sparse.linalg import aslinearoperator
+import scipy.linalg
+import scipy.sparse.linalg
+from scipy.linalg._decomp_qr import qr
 from scipy.sparse._sputils import is_pydata_spmatrix
+from scipy.sparse.linalg import aslinearoperator
+from scipy.sparse.linalg._interface import IdentityOperator
+from scipy.sparse.linalg._onenormest import onenormest
 
 __all__ = ['expm_multiply']
 
@@ -38,6 +42,53 @@ def _trace(A):
         return A.trace()
 
 
+def traceest(A, m3, seed=None):
+    """Estimate `np.trace(A)` using `3*m3` matrix-vector products.
+
+    The result is not deterministic.
+
+    Parameters
+    ----------
+    A : LinearOperator
+        Linear operator whose trace will be estimated. Has to be square.
+    m3 : int
+        Number of matrix-vector products divided by 3 used to estimate the
+        trace.
+    seed : optional
+        Seed for `numpy.random.default_rng`.
+        Can be provided to obtain deterministic results.
+
+    Returns
+    -------
+    trace : LinearOperator.dtype
+        Estimate of the trace
+
+    Notes
+    -----
+    This is the Hutch++ algorithm given in [1]_.
+
+    References
+    ----------
+    .. [1] Meyer, Raphael A., Cameron Musco, Christopher Musco, and David P.
+       Woodruff. "Hutch++: Optimal Stochastic Trace Estimation." In Symposium
+       on Simplicity in Algorithms (SOSA), pp. 142-155. Society for Industrial
+       and Applied Mathematics, 2021
+       https://doi.org/10.1137/1.9781611976496.16
+
+    """
+    rng = np.random.default_rng(seed)
+    if len(A.shape) != 2 or A.shape[-1] != A.shape[-2]:
+        raise ValueError("Expected A to be like a square matrix.")
+    n = A.shape[-1]
+    S = rng.choice([-1.0, +1.0], [n, m3])
+    Q, _ = qr(A.matmat(S), overwrite_a=True, mode='economic')
+    trQAQ = np.trace(Q.conj().T @ A.matmat(Q))
+    G = rng.choice([-1, +1], [n, m3])
+    right = G - Q@(Q.conj().T @ G)
+    trGAG = np.trace(right.conj().T @ A.matmat(right))
+    return trQAQ + trGAG/m3
+
+
 def _ident_like(A):
     # A compatibility function which should eventually disappear.
     if scipy.sparse.isspmatrix(A):
@@ -46,11 +97,14 @@ def _ident_like(A):
     elif is_pydata_spmatrix(A):
         import sparse
         return sparse.eye(A.shape[0], A.shape[1], dtype=A.dtype)
+    elif isinstance(A, scipy.sparse.linalg.LinearOperator):
+        return IdentityOperator(A.shape, dtype=A.dtype)
     else:
         return np.eye(A.shape[0], A.shape[1], dtype=A.dtype)
 
 
-def expm_multiply(A, B, start=None, stop=None, num=None, endpoint=None):
+def expm_multiply(A, B, start=None, stop=None, num=None,
+                  endpoint=None, traceA=None):
     """
     Compute the action of the matrix exponential of A on B.
 
@@ -71,12 +125,25 @@ def expm_multiply(A, B, start=None, stop=None, num=None, endpoint=None):
         Number of time points to use.
     endpoint : bool, optional
         If True, `stop` is the last time point.  Otherwise, it is not included.
+    traceA : scalar, optional
+        Trace of `A`. If not given the trace is estimated for linear operators,
+        or calculated exactly for sparse matrices. It is used to precondition
+        `A`, thus an approximate trace is acceptable.
+        For linear operators, `traceA` should be provided to ensure performance
+        as the estimation is not guaranteed to be reliable for all cases.
+
+        .. versionadded: 1.9.0
 
     Returns
     -------
     expm_A_B : ndarray
          The result of the action :math:`e^{t_k A} B`.
 
+    Warns
+    -----
+    UserWarning
+        If `A` is a linear operator and ``traceA==None`` (default).
+
     Notes
     -----
     The optional arguments defining the sequence of evenly spaced time points
@@ -133,13 +200,14 @@ def expm_multiply(A, B, start=None, stop=None, num=None, endpoint=None):
     array([ 2.71828183,  1.        ])
     """
     if all(arg is None for arg in (start, stop, num, endpoint)):
-        X = _expm_multiply_simple(A, B)
+        X = _expm_multiply_simple(A, B, traceA=traceA)
     else:
-        X, status = _expm_multiply_interval(A, B, start, stop, num, endpoint)
+        X, status = _expm_multiply_interval(A, B, start, stop, num,
+                                            endpoint, traceA=traceA)
     return X
 
 
-def _expm_multiply_simple(A, B, t=1.0, balance=False):
+def _expm_multiply_simple(A, B, t=1.0, traceA=None, balance=False):
     """
     Compute the action of the matrix exponential at a single time point.
 
@@ -151,6 +219,10 @@ def _expm_multiply_simple(A, B, t=1.0, balance=False):
         The matrix to be multiplied by the matrix exponential of A.
     t : float
         A time point.
+    traceA : scalar, optional
+        Trace of `A`. If not given the trace is estimated for linear operators,
+        or calculated exactly for sparse matrices. It is used to precondition
+        `A`, thus an approximate trace is acceptable
     balance : bool
         Indicates whether or not to apply balancing.
 
@@ -172,6 +244,7 @@ def _expm_multiply_simple(A, B, t=1.0, balance=False):
         raise ValueError('shapes of matrices A {} and B {} are incompatible'
                          .format(A.shape, B.shape))
     ident = _ident_like(A)
+    is_linear_operator = isinstance(A, scipy.sparse.linalg.LinearOperator)
     n = A.shape[0]
     if len(B.shape) == 1:
         n0 = 1
@@ -181,9 +254,16 @@ def _expm_multiply_simple(A, B, t=1.0, balance=False):
         raise ValueError('expected B to be like a matrix or a vector')
     u_d = 2**-53
     tol = u_d
-    mu = _trace(A) / float(n)
+    if traceA is None:
+        if is_linear_operator:
+            warn("Trace of LinearOperator not available, it will be estimated."
+                 " Provide `traceA` to ensure performance.", stacklevel=3)
+        # m3=1 is bit arbitrary choice, a more accurate trace (larger m3) might
+        # speed up exponential calculation, but trace estimation is more costly
+        traceA = traceest(A, m3=1) if is_linear_operator else _trace(A)
+    mu = traceA / float(n)
     A = A - mu * ident
-    A_1_norm = _exact_1_norm(A)
+    A_1_norm = onenormest(A) if is_linear_operator else _exact_1_norm(A)
     if t*A_1_norm == 0:
         m_star, s = 0, 1
     else:
@@ -508,8 +588,9 @@ def _condition_3_13(A_1_norm, n0, m_max, ell):
     return A_1_norm <= a * b
 
 
-def _expm_multiply_interval(A, B, start=None, stop=None,
-        num=None, endpoint=None, balance=False, status_only=False):
+def _expm_multiply_interval(A, B, start=None, stop=None, num=None,
+                            endpoint=None, traceA=None, balance=False,
+                            status_only=False):
     """
     Compute the action of the matrix exponential at multiple time points.
 
@@ -528,6 +609,10 @@ def _expm_multiply_interval(A, B, start=None, stop=None,
         Note that the step size changes when `endpoint` is False.
     num : int, optional
         Number of time points to use.
+    traceA : scalar, optional
+        Trace of `A`. If not given the trace is estimated for linear operators,
+        or calculated exactly for sparse matrices. It is used to precondition
+        `A`, thus an approximate trace is acceptable
     endpoint : bool, optional
         If True, `stop` is the last time point. Otherwise, it is not included.
     balance : bool
@@ -558,6 +643,7 @@ def _expm_multiply_interval(A, B, start=None, stop=None,
         raise ValueError('shapes of matrices A {} and B {} are incompatible'
                          .format(A.shape, B.shape))
     ident = _ident_like(A)
+    is_linear_operator = isinstance(A, scipy.sparse.linalg.LinearOperator)
     n = A.shape[0]
     if len(B.shape) == 1:
         n0 = 1
@@ -567,7 +653,15 @@ def _expm_multiply_interval(A, B, start=None, stop=None,
         raise ValueError('expected B to be like a matrix or a vector')
     u_d = 2**-53
     tol = u_d
-    mu = _trace(A) / float(n)
+    if traceA is None:
+        if is_linear_operator:
+            warn("Trace of LinearOperator not available, it will be estimated."
+                 " Provide `traceA` to ensure performance.", stacklevel=3)
+        # m3=5 is bit arbitrary choice, a more accurate trace (larger m3) might
+        # speed up exponential calculation, but trace estimation is also costly
+        # an educated guess would need to consider the number of time points
+        traceA = traceest(A, m3=5) if is_linear_operator else _trace(A)
+    mu = traceA / float(n)
 
     # Get the linspace samples, attempting to preserve the linspace defaults.
     linspace_kwargs = {'retstep': True}
@@ -593,7 +687,7 @@ def _expm_multiply_interval(A, B, start=None, stop=None,
     X = np.empty(X_shape, dtype=np.result_type(A.dtype, B.dtype, float))
     t = t_q - t_0
     A = A - mu * ident
-    A_1_norm = _exact_1_norm(A)
+    A_1_norm = onenormest(A) if is_linear_operator else _exact_1_norm(A)
     ell = 2
     norm_info = LazyOperatorNormInfo(t*A, A_1_norm=t*A_1_norm, ell=ell)
     if t*A_1_norm == 0:


Test patch:
diff --git a/scipy/sparse/linalg/tests/test_expm_multiply.py b/scipy/sparse/linalg/tests/test_expm_multiply.py
--- a/scipy/sparse/linalg/tests/test_expm_multiply.py
+++ b/scipy/sparse/linalg/tests/test_expm_multiply.py
@@ -1,10 +1,13 @@
-"""Test functions for the sparse.linalg._expm_multiply module
-"""
+"""Test functions for the sparse.linalg._expm_multiply module."""
+from functools import partial
+from itertools import product
 
 import numpy as np
+import pytest
 from numpy.testing import (assert_allclose, assert_, assert_equal,
                            suppress_warnings)
 from scipy.sparse import SparseEfficiencyWarning
+from scipy.sparse.linalg import aslinearoperator
 import scipy.linalg
 from scipy.sparse.linalg import expm as sp_expm
 from scipy.sparse.linalg._expm_multiply import (_theta, _compute_p_max,
@@ -12,6 +15,28 @@
         _expm_multiply_interval)
 
 
+IMPRECISE = {np.single, np.csingle}
+REAL_DTYPES = {np.intc, np.int_, np.longlong,
+               np.single, np.double, np.longdouble}
+COMPLEX_DTYPES = {np.csingle, np.cdouble, np.clongdouble}
+# use sorted tuple to ensure fixed order of tests
+DTYPES = tuple(sorted(REAL_DTYPES ^ COMPLEX_DTYPES, key=str))
+
+
+def estimated(func):
+    """If trace is estimated, it should warn.
+
+    We warn that estimation of trace might impact performance.
+    All result have to be correct nevertheless!
+
+    """
+    def wrapped(*args, **kwds):
+        with pytest.warns(UserWarning,
+                          match="Trace of LinearOperator not available"):
+            return func(*args, **kwds)
+    return wrapped
+
+
 def less_than_or_close(a, b):
     return np.allclose(a, b) or (a < b)
 
@@ -66,6 +91,11 @@ def test_expm_multiply(self):
             observed = expm_multiply(A, B)
             expected = np.dot(sp_expm(A), B)
             assert_allclose(observed, expected)
+            observed = estimated(expm_multiply)(aslinearoperator(A), B)
+            assert_allclose(observed, expected)
+            traceA = np.trace(A)
+            observed = expm_multiply(aslinearoperator(A), B, traceA=traceA)
+            assert_allclose(observed, expected)
 
     def test_matrix_vector_multiply(self):
         np.random.seed(1234)
@@ -77,20 +107,25 @@ def test_matrix_vector_multiply(self):
             observed = expm_multiply(A, v)
             expected = np.dot(sp_expm(A), v)
             assert_allclose(observed, expected)
+            observed = estimated(expm_multiply)(aslinearoperator(A), v)
+            assert_allclose(observed, expected)
 
     def test_scaled_expm_multiply(self):
         np.random.seed(1234)
         n = 40
         k = 3
         nsamples = 10
-        for i in range(nsamples):
-            for t in (0.2, 1.0, 1.5):
-                with np.errstate(invalid='ignore'):
-                    A = scipy.linalg.inv(np.random.randn(n, n))
-                    B = np.random.randn(n, k)
-                    observed = _expm_multiply_simple(A, B, t=t)
-                    expected = np.dot(sp_expm(t*A), B)
-                    assert_allclose(observed, expected)
+        for i, t in product(range(nsamples), [0.2, 1.0, 1.5]):
+            with np.errstate(invalid='ignore'):
+                A = scipy.linalg.inv(np.random.randn(n, n))
+                B = np.random.randn(n, k)
+                observed = _expm_multiply_simple(A, B, t=t)
+                expected = np.dot(sp_expm(t*A), B)
+                assert_allclose(observed, expected)
+                observed = estimated(_expm_multiply_simple)(
+                    aslinearoperator(A), B, t=t
+                )
+                assert_allclose(observed, expected)
 
     def test_scaled_expm_multiply_single_timepoint(self):
         np.random.seed(1234)
@@ -102,6 +137,10 @@ def test_scaled_expm_multiply_single_timepoint(self):
         observed = _expm_multiply_simple(A, B, t=t)
         expected = sp_expm(t*A).dot(B)
         assert_allclose(observed, expected)
+        observed = estimated(_expm_multiply_simple)(
+            aslinearoperator(A), B, t=t
+        )
+        assert_allclose(observed, expected)
 
     def test_sparse_expm_multiply(self):
         np.random.seed(1234)
@@ -120,6 +159,8 @@ def test_sparse_expm_multiply(self):
                            " CSC matrix format")
                 expected = sp_expm(A).dot(B)
             assert_allclose(observed, expected)
+            observed = estimated(expm_multiply)(aslinearoperator(A), B)
+            assert_allclose(observed, expected)
 
     def test_complex(self):
         A = np.array([
@@ -131,6 +172,8 @@ def test_complex(self):
             1j * np.exp(1j) + 1j * (1j*np.cos(1) - np.sin(1)),
             1j * np.exp(1j)], dtype=complex)
         assert_allclose(observed, expected)
+        observed = estimated(expm_multiply)(aslinearoperator(A), B)
+        assert_allclose(observed, expected)
 
 
 class TestExpmActionInterval:
@@ -147,50 +190,59 @@ def test_sparse_expm_multiply_interval(self):
             B = np.random.randn(n, k)
             v = np.random.randn(n)
             for target in (B, v):
-                X = expm_multiply(A, target,
-                        start=start, stop=stop, num=num, endpoint=endpoint)
+                X = expm_multiply(A, target, start=start, stop=stop,
+                                  num=num, endpoint=endpoint)
                 samples = np.linspace(start=start, stop=stop,
-                        num=num, endpoint=endpoint)
+                                      num=num, endpoint=endpoint)
                 with suppress_warnings() as sup:
                     sup.filter(SparseEfficiencyWarning,
                                "splu converted its input to CSC format")
                     sup.filter(SparseEfficiencyWarning,
-                               "spsolve is more efficient when sparse b is in the CSC matrix format")
+                               "spsolve is more efficient when sparse b is in"
+                               " the CSC matrix format")
                     for solution, t in zip(X, samples):
                         assert_allclose(solution, sp_expm(t*A).dot(target))
 
     def test_expm_multiply_interval_vector(self):
         np.random.seed(1234)
-        start = 0.1
-        stop = 3.2
-        endpoint = True
-        for num in (14, 13, 2):
-            for n in (1, 2, 5, 20, 40):
-                A = scipy.linalg.inv(np.random.randn(n, n))
-                v = np.random.randn(n)
-                X = expm_multiply(A, v,
-                        start=start, stop=stop, num=num, endpoint=endpoint)
-                samples = np.linspace(start=start, stop=stop,
-                        num=num, endpoint=endpoint)
-                for solution, t in zip(X, samples):
-                    assert_allclose(solution, sp_expm(t*A).dot(v))
+        interval = {'start': 0.1, 'stop': 3.2, 'endpoint': True}
+        for num, n in product([14, 13, 2], [1, 2, 5, 20, 40]):
+            A = scipy.linalg.inv(np.random.randn(n, n))
+            v = np.random.randn(n)
+            samples = np.linspace(num=num, **interval)
+            X = expm_multiply(A, v, num=num, **interval)
+            for solution, t in zip(X, samples):
+                assert_allclose(solution, sp_expm(t*A).dot(v))
+            # test for linear operator with unknown trace -> estimate trace
+            Xguess = estimated(expm_multiply)(aslinearoperator(A), v,
+                                              num=num, **interval)
+            # test for linear operator with given trace
+            Xgiven = expm_multiply(aslinearoperator(A), v, num=num, **interval,
+                                   traceA=np.trace(A))
+            # test robustness for linear operator with wrong trace
+            Xwrong = expm_multiply(aslinearoperator(A), v, num=num, **interval,
+                                   traceA=np.trace(A)*5)
+            for sol_guess, sol_given, sol_wrong, t in zip(Xguess, Xgiven,
+                                                          Xwrong, samples):
+                correct = sp_expm(t*A).dot(v)
+                assert_allclose(sol_guess, correct)
+                assert_allclose(sol_given, correct)
+                assert_allclose(sol_wrong, correct)
 
     def test_expm_multiply_interval_matrix(self):
         np.random.seed(1234)
-        start = 0.1
-        stop = 3.2
-        endpoint = True
-        for num in (14, 13, 2):
-            for n in (1, 2, 5, 20, 40):
-                for k in (1, 2):
-                    A = scipy.linalg.inv(np.random.randn(n, n))
-                    B = np.random.randn(n, k)
-                    X = expm_multiply(A, B,
-                            start=start, stop=stop, num=num, endpoint=endpoint)
-                    samples = np.linspace(start=start, stop=stop,
-                            num=num, endpoint=endpoint)
-                    for solution, t in zip(X, samples):
-                        assert_allclose(solution, sp_expm(t*A).dot(B))
+        interval = {'start': 0.1, 'stop': 3.2, 'endpoint': True}
+        for num, n, k in product([14, 13, 2], [1, 2, 5, 20, 40], [1, 2]):
+            A = scipy.linalg.inv(np.random.randn(n, n))
+            B = np.random.randn(n, k)
+            samples = np.linspace(num=num, **interval)
+            X = expm_multiply(A, B, num=num, **interval)
+            for solution, t in zip(X, samples):
+                assert_allclose(solution, sp_expm(t*A).dot(B))
+            X = estimated(expm_multiply)(aslinearoperator(A), B, num=num,
+                                         **interval)
+            for solution, t in zip(X, samples):
+                assert_allclose(solution, sp_expm(t*A).dot(B))
 
     def test_sparse_expm_multiply_interval_dtypes(self):
         # Test A & B int
@@ -249,3 +301,45 @@ def _help_test_specific_expm_interval_status(self, target_status):
         if not nsuccesses:
             msg = 'failed to find a status-' + str(target_status) + ' interval'
             raise Exception(msg)
+
+
+@pytest.mark.parametrize("dtype_a", DTYPES)
+@pytest.mark.parametrize("dtype_b", DTYPES)
+@pytest.mark.parametrize("b_is_matrix", [False, True])
+def test_expm_multiply_dtype(dtype_a, dtype_b, b_is_matrix):
+    """Make sure `expm_multiply` handles all numerical dtypes correctly."""
+    assert_allclose_ = (partial(assert_allclose, rtol=1e-3, atol=1e-5)
+                        if {dtype_a, dtype_b} & IMPRECISE else assert_allclose)
+    rng = np.random.default_rng(1234)
+    # test data
+    n = 7
+    b_shape = (n, 3) if b_is_matrix else (n, )
+    if dtype_a in REAL_DTYPES:
+        A = scipy.linalg.inv(rng.random([n, n])).astype(dtype_a)
+    else:
+        A = scipy.linalg.inv(
+            rng.random([n, n]) + 1j*rng.random([n, n])
+        ).astype(dtype_a)
+    if dtype_b in REAL_DTYPES:
+        B = (2*rng.random(b_shape)).astype(dtype_b)
+    else:
+        B = (rng.random(b_shape) + 1j*rng.random(b_shape)).astype(dtype_b)
+
+    # single application
+    sol_mat = expm_multiply(A, B)
+    sol_op = estimated(expm_multiply)(aslinearoperator(A), B)
+    direct_sol = np.dot(sp_expm(A), B)
+    assert_allclose_(sol_mat, direct_sol)
+    assert_allclose_(sol_op, direct_sol)
+    sol_op = expm_multiply(aslinearoperator(A), B, traceA=np.trace(A))
+    assert_allclose_(sol_op, direct_sol)
+
+    # for time points
+    interval = {'start': 0.1, 'stop': 3.2, 'num': 13, 'endpoint': True}
+    samples = np.linspace(**interval)
+    X_mat = expm_multiply(A, B, **interval)
+    X_op = estimated(expm_multiply)(aslinearoperator(A), B, **interval)
+    for sol_mat, sol_op, t in zip(X_mat, X_op, samples):
+        direct_sol = sp_expm(t*A).dot(B)
+        assert_allclose_(sol_mat, direct_sol)
+        assert_allclose_(sol_op, direct_sol)