Refactor TestSvd and TestNormTests to handle complex and batched cases

Author: jessegrabowski

pytensor/tensor/nlinalg.py

@@ -559,7 +559,7 @@ def make_node(self, x):
         assert x.ndim == 2, "The input of svd function should be a matrix."
 
         in_dtype = x.type.numpy_dtype
-        out_dtype = np.dtype(f"f{in_dtype.itemsize}")
+        out_dtype = in_dtype
 
         s = vector(dtype=out_dtype)
 

tests/tensor/test_nlinalg.py

@@ -1,3 +1,5 @@
+from functools import partial
+
 import numpy as np
 import pytest
 from numpy.testing import assert_array_almost_equal
@@ -148,41 +150,51 @@ def test_qr_modes():
 
 class TestSvd(utt.InferShapeTester):
     op_class = SVD
-    dtype = "float32"
 
     def setup_method(self):
         super().setup_method()
         self.rng = np.random.default_rng(utt.fetch_seed())
-        self.A = matrix(dtype=self.dtype)
+        self.A = matrix(dtype=config.floatX)
         self.op = svd
 
-    def test_svd(self):
-        A = matrix("A", dtype=self.dtype)
-        U, S, VT = svd(A)
-        fn = function([A], [U, S, VT])
-        a = self.rng.random((4, 4)).astype(self.dtype)
-        n_u, n_s, n_vt = np.linalg.svd(a)
-        t_u, t_s, t_vt = fn(a)
-
-        assert _allclose(n_u, t_u)
-        assert _allclose(n_s, t_s)
-        assert _allclose(n_vt, t_vt)
-
-        fn = function([A], svd(A, compute_uv=False))
-        t_s = fn(a)
-        assert _allclose(n_s, t_s)
-
-        A = tensor3("A", dtype=self.dtype)
-        U, S, VT = svd(A)
-        fn = function([A], [U, S, VT])
-        a = self.rng.random((10, 4, 4)).astype(self.dtype)
-        t_u, t_s, t_vt = fn(a)
-        n_u, n_s, n_vt = np.vectorize(
-            np.linalg.svd, signature="(i,j)->(i,k),(k),(k,j)"
-        )(a)
-        assert _allclose(n_u, t_u)
-        assert _allclose(n_s, t_s)
-        assert _allclose(n_vt, t_vt)
+    @pytest.mark.parametrize(
+        "compute_uv", [True, False], ids=["compute_uv=True", "compute_uv=False"]
+    )
+    @pytest.mark.parametrize(
+        "batched", [True, False], ids=["batched=True", "batched=False"]
+    )
+    @pytest.mark.parametrize(
+        "test_imag", [True, False], ids=["test_imag=True", "test_imag=False"]
+    )
+    def test_svd(self, compute_uv, batched, test_imag):
+        dtype = config.floatX
+        if test_imag:
+            dtype = "complex128" if dtype.endswith("64") else "complex64"
+
+        if batched:
+            A = tensor3("A", dtype=dtype)
+            size = (10, 4, 4)
+        else:
+            A = matrix("A", dtype=dtype)
+            size = (4, 4)
+        a = self.rng.random(size).astype(dtype)
+
+        outputs = svd(A, compute_uv=compute_uv, full_matrices=False)
+        outputs = outputs if isinstance(outputs, list) else [outputs]
+        fn = function(inputs=[A], outputs=outputs)
+
+        np_fn = np.vectorize(
+            partial(np.linalg.svd, compute_uv=compute_uv, full_matrices=False),
+            signature=outputs[0].owner.op.core_op.gufunc_signature,
+        )
+
+        np_outputs = np_fn(a)
+        pt_outputs = fn(a)
+
+        np_outputs = np_outputs if isinstance(np_outputs, tuple) else [np_outputs]
+
+        for np_val, pt_val in zip(np_outputs, pt_outputs):
+            assert _allclose(np_val, pt_val)
 
     def test_svd_infer_shape(self):
         self.validate_shape((4, 4), full_matrices=True, compute_uv=True)
@@ -193,7 +205,7 @@ def test_svd_infer_shape(self):
 
     def validate_shape(self, shape, compute_uv=True, full_matrices=True):
         A = self.A
-        A_v = self.rng.random(shape).astype(self.dtype)
+        A_v = self.rng.random(shape).astype(config.floatX)
         outputs = self.op(A, full_matrices=full_matrices, compute_uv=compute_uv)
         if not compute_uv:
             outputs = [outputs]
@@ -465,9 +477,9 @@ def test_non_tensorial_input(self):
         [None, np.inf, -np.inf, 1, -1, 2, -2],
         ids=["None", "inf", "-inf", "1", "-1", "2", "-2"],
     )
-    @pytest.mark.parametrize("core_dims", [(4,), (4, 4)], ids=["vector", "matrix"])
+    @pytest.mark.parametrize("core_dims", [(4,), (4, 3)], ids=["vector", "matrix"])
     @pytest.mark.parametrize("batch_dims", [(), (2,)], ids=["no_batch", "batch"])
-    @pytest.mark.parametrize("test_imag", [True, False], ids=["real", "complex"])
+    @pytest.mark.parametrize("test_imag", [True, False], ids=["complex", "real"])
     def test_numpy_compare(
         self,
         ord: float,
@@ -481,6 +493,8 @@ def test_numpy_compare(
         has_batch = len(batch_dims) > 0
         if ord in [np.inf, -np.inf] and not is_matrix:
             pytest.skip("Infinity norm not defined for vectors")
+        if test_imag and is_matrix and ord == -2:
+            pytest.skip("Complex matrices not supported")
         if has_batch and not is_matrix:
             # Handle batched vectors by row-normalizing a matrix
             axis = (-1,)
@@ -491,8 +505,8 @@ def test_numpy_compare(
             x_real, x_imag = rng.standard_normal((2,) + batch_dims + core_dims).astype(
                 config.floatX
             )
-            dtype = "complex64" if config.floatX.endswith("64") else "complex32"
-            X = x_real.astype(dtype) + 1j * x_imag.astype(dtype)
+            dtype = "complex128" if config.floatX.endswith("64") else "complex64"
+            X = (x_real + 1j * x_imag).astype(dtype)
         else:
             X = rng.standard_normal(batch_dims + core_dims).astype(config.floatX)
 
@@ -505,6 +519,7 @@ def test_numpy_compare(
 
         pt_norm = norm(X, ord=ord, axis=axis, keepdims=keepdims)
         f = function([], pt_norm, mode="FAST_COMPILE")
+
         utt.assert_allclose(np_norm, f())