Remove SquareDiagonal Op and replace it with a square_diagonal function

Author: tomicapretto

pytensor/sparse/basic.py

@@ -1445,53 +1445,13 @@ def infer_shape(self, fgraph, nodes, shapes):
 diag = Diag()
 
 
-class SquareDiagonal(Op):
-    """Produce a square sparse (csc) matrix with a diagonal given by a dense vector.
-
-    Notes
-    -----
-    The grad implemented is regular, i.e. not structured.
-
-    """
-
-    __props__ = ()
-
-    def make_node(self, diag):
-        """
-
-        Parameters
-        ----------
-        x
-            Dense vector for the diagonal.
-
-        """
-        diag = ptb.as_tensor_variable(diag)
-        if diag.type.ndim != 1:
-            raise TypeError("data argument must be a vector", diag.type)
-
-        return Apply(self, [diag], [SparseTensorType(dtype=diag.dtype, format="csc")()])
-
-    def perform(self, node, inputs, outputs):
-        (z,) = outputs
-        diag = inputs[0]
-
-        N = len(diag)
-        data = diag[:N]
-        indices = list(range(N))
-        indptr = list(range(N + 1))
-        tup = (data, indices, indptr)
-
-        z[0] = scipy.sparse.csc_matrix(tup, copy=True)
-
-    def grad(self, inputs, gout):
-        (gz,) = gout
-        return [diag(gz)]
-
-    def infer_shape(self, fgraph, nodes, shapes):
-        return [(shapes[0][0], shapes[0][0])]
-
-
-square_diagonal = SquareDiagonal()
+def square_diagonal(diag):
+    """Produce a square sparse (csc) matrix with a diagonal given by a dense vector."""
+    n = diag.shape[0]
+    data = ptb.as_tensor_variable(diag)
+    indices = ptb.arange(n, dtype=np.int32)
+    indptr = ptb.arange(n + 1, dtype=np.int32)
+    return CSC(data, indices, indptr, ptb.as_tensor((n, n)))
 
 
 class EnsureSortedIndices(Op):

tests/sparse/test_basic.py

@@ -29,7 +29,6 @@
     Remove0,
     SparseFromDense,
     SparseTensorType,
-    SquareDiagonal,
     Transpose,
     VStack,
     _is_sparse,
@@ -966,42 +965,33 @@ def test_grad(self):
             verify_grad_sparse(self.op, data, structured=False)
 
 
-class TestSquareDiagonal(utt.InferShapeTester):
-    def setup_method(self):
-        super().setup_method()
-        self.op_class = SquareDiagonal
-        self.op = square_diagonal
-
+class TestSquareDiagonal:
     def test_op(self):
-        for format in sparse.sparse_formats:
-            for size in range(5, 9):
-                variable = [vector()]
-                data = [np.random.random(size).astype(config.floatX)]
-
-                f = pytensor.function(variable, self.op(*variable))
-                tested = f(*data).toarray()
-
-                expected = np.diag(*data)
-                utt.assert_allclose(expected, tested)
-                assert tested.dtype == expected.dtype
-                assert tested.shape == expected.shape
-
-    def test_infer_shape(self):
-        for format in sparse.sparse_formats:
-            for size in range(5, 9):
-                variable = [vector()]
-                data = [np.random.random(size).astype(config.floatX)]
-
-                self._compile_and_check(
-                    variable, [self.op(*variable)], data, self.op_class
-                )
+        x = vector(dtype=config.floatX)
+        y = square_diagonal(x)
+        f = pytensor.function([x], y)
+
+        size = 11
+        values = np.random.random(size).astype(config.floatX)
+        tested = f(values)
+
+        assert tested.format == "csc"
+        utt.assert_allclose(tested.toarray(), np.diag(values))
+        assert tuple(tested.shape) == (values.size, values.size)
+        np.testing.assert_array_equal(
+            tested.indices, np.arange(values.size, dtype="int32")
+        )
+        np.testing.assert_array_equal(
+            tested.indptr, np.arange(values.size + 1, dtype="int32")
+        )
 
     def test_grad(self):
-        for format in sparse.sparse_formats:
-            for size in range(5, 9):
-                data = [np.random.random(size).astype(config.floatX)]
+        values = [np.random.random(13).astype(config.floatX)]
+        verify_grad_sparse(square_diagonal, values, structured=False)
 
-                verify_grad_sparse(self.op, data, structured=False)
+    def test_rejects_non_vector_input(self):
+        with pytest.raises(TypeError, match="data argument must be a vector"):
+            square_diagonal(matrix(dtype=config.floatX))
 
 
 class TestEnsureSortedIndices(utt.InferShapeTester):