Clean up pytensor.linalg.expm and related tests

Author: Jesse Grabowski

pytensor/tensor/slinalg.py

@@ -1304,82 +1304,63 @@ def eigvalsh(a, b, lower=True):
 class Expm(Op):
     """
     Compute the matrix exponential of a square array.
-
     """
 
     __props__ = ()
+    gufunc_signature = "(m,m)->(m,m)"
 
     def make_node(self, A):
         A = as_tensor_variable(A)
         assert A.ndim == 2
-        expm = matrix(dtype=A.dtype)
-        return Apply(
-            self,
-            [
-                A,
-            ],
-            [
-                expm,
-            ],
-        )
+
+        expm = matrix(dtype=A.dtype, shape=A.type.shape)
+
+        return Apply(self, [A], [expm])
 
     def perform(self, node, inputs, outputs):
         (A,) = inputs
         (expm,) = outputs
         expm[0] = scipy_linalg.expm(A)
 
-    def grad(self, inputs, outputs):
+    def L_op(self, inputs, outputs, output_grads):
+        # Kalbfleisch and Lawless, J. Am. Stat. Assoc. 80 (1985) Equation 3.4
+        # Kind of... You need to do some algebra from there to arrive at
+        # this expression.
         (A,) = inputs
-        (g_out,) = outputs
-        return [ExpmGrad()(A, g_out)]
-
-    def infer_shape(self, fgraph, node, shapes):
-        return [shapes[0]]
+        (_,) = outputs  # Outputs not used; included for signature consistency only
+        (A_bar,) = output_grads
 
+        w, V = pt.linalg.eig(A, return_components=True)
 
-class ExpmGrad(Op):
-    """
-    Gradient of the matrix exponential of a square array.
+        w = w[0] + 1j * w[1]
+        V = V[0] + 1j * V[1]
 
-    """
+        exp_w = pt.exp(w)
+        numer = pt.sub.outer(exp_w, exp_w)
+        denom = pt.sub.outer(w, w)
 
-    __props__ = ()
+        # When w_i ≈ w_j, we have a removable singularity in the expression for X, because
+        # lim b->a (e^a - e^b) / (a - b) = e^a (derivation left for the motivated reader)
+        X = pt.where(pt.abs(denom) < 1e-8, exp_w, numer / denom)
 
-    def make_node(self, A, gw):
-        A = as_tensor_variable(A)
-        assert A.ndim == 2
-        out = matrix(dtype=A.dtype)
-        return Apply(
-            self,
-            [A, gw],
-            [
-                out,
-            ],
-        )
+        diag_idx = pt.arange(w.shape[0])
+        X = X[..., diag_idx, diag_idx].set(exp_w)
 
-    def infer_shape(self, fgraph, node, shapes):
-        return [shapes[0]]
+        inner = solve(V, A_bar.T @ V).T
+        result = solve(V.T, inner * X) @ V.T
 
-    def perform(self, node, inputs, outputs):
-        # Kalbfleisch and Lawless, J. Am. Stat. Assoc. 80 (1985) Equation 3.4
-        # Kind of... You need to do some algebra from there to arrive at
-        # this expression.
-        (A, gA) = inputs
-        (out,) = outputs
-        w, V = scipy_linalg.eig(A, right=True)
-        U = scipy_linalg.inv(V).T
+        # At this point, result is always a complex dtype. If the input was real, the output should be
+        # real as well (and all the imaginary parts are numerical noise)
+        if A.dtype not in ("complex64", "complex128"):
+            return [result.real]
 
-        exp_w = np.exp(w)
-        X = np.subtract.outer(exp_w, exp_w) / np.subtract.outer(w, w)
-        np.fill_diagonal(X, exp_w)
-        Y = U.dot(V.T.dot(gA).dot(U) * X).dot(V.T)
+        return [result]
 
-        with warnings.catch_warnings():
-            warnings.simplefilter("ignore", ComplexWarning)
-            out[0] = Y.astype(A.dtype)
+    def infer_shape(self, fgraph, node, shapes):
+        return [shapes[0]]
 
 
-expm = Expm()
+expm = Blockwise(Expm())
 
 
 class SolveContinuousLyapunov(Op):

tests/tensor/test_slinalg.py

@@ -880,35 +880,26 @@ def test_expm():
     np.testing.assert_array_almost_equal(val, ref)
 
 
-def test_expm_grad_1():
-    # with symmetric matrix (real eigenvectors)
-    rng = np.random.default_rng(utt.fetch_seed())
-    # Always test in float64 for better numerical stability.
-    A = rng.standard_normal((5, 5))
-    A = A + A.T
-
-    utt.verify_grad(expm, [A], rng=rng)
-
-
-def test_expm_grad_2():
-    # with non-symmetric matrix with real eigenspecta
-    rng = np.random.default_rng(utt.fetch_seed())
-    # Always test in float64 for better numerical stability.
-    A = rng.standard_normal((5, 5))
-    w = rng.standard_normal(5) ** 2
-    A = (np.diag(w**0.5)).dot(A + A.T).dot(np.diag(w ** (-0.5)))
-    assert not np.allclose(A, A.T)
-
-    utt.verify_grad(expm, [A], rng=rng)
-
-
-def test_expm_grad_3():
-    # with non-symmetric matrix (complex eigenvectors)
-    rng = np.random.default_rng(utt.fetch_seed())
-    # Always test in float64 for better numerical stability.
-    A = rng.standard_normal((5, 5))
-
-    utt.verify_grad(expm, [A], rng=rng)
+@pytest.mark.parametrize(
+    "mode", ["symmetric", "nonsymmetric_real_eig", "nonsymmetric_complex_eig"][-1:]
+)
+def test_expm_grad(mode):
+    rng = np.random.default_rng()
+
+    match mode:
+        case "symmetric":
+            A = rng.standard_normal((5, 5))
+            A = A + A.T
+        case "nonsymmetric_real_eig":
+            A = rng.standard_normal((5, 5))
+            w = rng.standard_normal(5) ** 2
+            A = (np.diag(w**0.5)).dot(A + A.T).dot(np.diag(w ** (-0.5)))
+        case "nonsymmetric_complex_eig":
+            A = rng.standard_normal((5, 5))
+        case _:
+            raise ValueError(f"Invalid mode: {mode}")
+
+    utt.verify_grad(expm, [A], rng=rng, abs_tol=1e-5, rel_tol=1e-5)
 
 
 def recover_Q(A, X, continuous=True):