Add return_components argument to pt.linalg.eig

This argument allows the user to get back both the real and imaginary parts of the eigenvalues, if required.

Author: Jesse Grabowski

pytensor/tensor/nlinalg.py

@@ -16,7 +16,15 @@
 from pytensor.tensor import math as ptm
 from pytensor.tensor.basic import as_tensor_variable, diagonal
 from pytensor.tensor.blockwise import Blockwise
-from pytensor.tensor.type import Variable, dvector, lscalar, matrix, scalar, vector
+from pytensor.tensor.type import (
+    Variable,
+    dvector,
+    lscalar,
+    matrix,
+    scalar,
+    tensor3,
+    vector,
+)
 
 
 class MatrixPinv(Op):
@@ -297,31 +305,68 @@ def slogdet(x: TensorLike) -> tuple[ptb.TensorVariable, ptb.TensorVariable]:
 class Eig(Op):
     """
     Compute the eigenvalues and right eigenvectors of a square array.
-
     """
 
-    __props__: tuple[str, ...] = ()
-    gufunc_signature = "(m,m)->(m),(m,m)"
+    __props__: tuple[str, ...] = ("return_components",)
     gufunc_spec = ("numpy.linalg.eig", 1, 2)
 
+    def __init__(self, return_components: bool = False):
+        self.return_components = return_components
+        if return_components:
+            signature = "(m,m)->(a,m),(a,m,m)"
+        else:
+            signature = "(m,m)->(m),(m,m)"
+        self.gufunc_signature = signature
+
     def make_node(self, x):
         x = as_tensor_variable(x)
         assert x.ndim == 2
-        w = vector(dtype=x.dtype)
-        v = matrix(dtype=x.dtype)
+
+        if self.return_components:
+            w = matrix(dtype=x.dtype)
+            v = tensor3(dtype=x.dtype)
+
+        else:
+            w = vector(dtype=x.dtype)
+            v = matrix(dtype=x.dtype)
+
         return Apply(self, [x], [w, v])
 
     def perform(self, node, inputs, outputs):
         (x,) = inputs
         (w, v) = outputs
-        w[0], v[0] = (z.astype(x.dtype) for z in np.linalg.eig(x))
+        if self.return_components:
+            w_res, v_res = np.linalg.eig(x)
+            w[0] = np.stack([w_res.real, w_res.imag], axis=0)
+            v[0] = np.stack([v_res.real, v_res.imag], axis=0)
+        else:
+            w[0], v[0] = (z.astype(x.dtype) for z in np.linalg.eig(x))
 
     def infer_shape(self, fgraph, node, shapes):
         n = shapes[0][0]
+        if self.return_components:
+            return [(2, n), (2, n, n)]
         return [(n,), (n, n)]
 
 
-eig = Blockwise(Eig())
+def eig(x: TensorLike, return_components: bool = False):
+    """
+    Return the eigenvalues and right eigenvectors of a square array.
+
+    Parameters
+    ----------
+    x: TensorLike
+        Square matrix, or array of such matrices
+    return_components: bool, optional
+        By default, only the real component of the eigenvalues and eigenvectors are returned, as Pytensor does not
+        allow the dtype of a variable to change during graph execution.
+
+        To circumvent this, if `return_components` is set to True, the real and imaginary *components* of the
+        eigenvalues are returned as a concatenated array of shape (2, n), where the first row contains the real parts
+        and the second row contains the imaginary parts. Similarly, the eigenvectors are returned as an array of shape
+        (2, n, n).
+    """
+    return Blockwise(Eig(return_components=return_components))(x)
 
 
 class Eigh(Eig):

tests/tensor/test_nlinalg.py

@@ -394,11 +394,12 @@ def test_trace():
 
 class TestEig(utt.InferShapeTester):
     op_class = Eig
-    op = eig
     dtype = "float64"
+    op = staticmethod(eig)
 
     def setup_method(self):
         super().setup_method()
+
         self.rng = np.random.default_rng(utt.fetch_seed())
         self.A = matrix(dtype=self.dtype)
         self.X = np.asarray(self.rng.random((5, 5)), dtype=self.dtype)
@@ -423,6 +424,23 @@ def test_eval(self):
         w, v = (e.eval({A: x}) for e in self.op(A))
         assert_array_almost_equal(np.dot(x, v), w * v)
 
+    def test_eval_return_components(self):
+        A = matrix(dtype=self.dtype)
+        A_val = self.rng.normal(size=(5, 5))
+        w, v = (e.eval({A: A_val}) for e in self.op(A, return_components=True))
+        assert w.shape == (2, 5)
+        assert v.shape == (2, 5, 5)
+
+        w = w[0] + 1j * w[1]
+        v = v[0] + 1j * v[1]
+
+        w_np, v_np = np.linalg.eig(A_val)
+
+        np.testing.assert_allclose(w, w_np)
+        np.testing.assert_allclose(v, v_np)
+
+        assert_array_almost_equal(A_val @ v, w * v)
+
 
 class TestEigh(TestEig):
     op = staticmethod(eigh)