[WIP] interface dense matmul support in SparseDiffPy

cvxpy/reductions/solvers/nlp_solvers/diff_engine/converters.py

@@ -32,7 +32,6 @@ def normalize_shape(shape):
     shape = tuple(shape)
     return (1,) * (2 - len(shape)) + shape
 
-
 def _chain_add(children):
     """Chain multiple children with binary adds: a + b + c -> add(add(a, b), c)."""
     result = children[0]
@@ -48,31 +47,49 @@ def _convert_matmul(expr, children):
     if left_arg.is_constant():
         A = left_arg.value
 
-        if not isinstance(A, sparse.csr_matrix):
-          A = sparse.csr_matrix(A)
-
-        return _diffengine.make_left_matmul(
-            children[1],
-            A.data.astype(np.float64),
-            A.indices.astype(np.int32),
-            A.indptr.astype(np.int32),
-            A.shape[0],
-            A.shape[1],
-        )
+        if sparse.issparse(A):
+            if not isinstance(A, sparse.csr_matrix):
+                A = sparse.csr_matrix(A)
+
+            return _diffengine.make_sparse_left_matmul(
+                children[1],
+                A.data.astype(np.float64, copy=False),
+                A.indices.astype(np.int32, copy=False),
+                A.indptr.astype(np.int32, copy=False),
+                A.shape[0],
+                A.shape[1],
+            )
+        else:
+            m, n = normalize_shape(A.shape)
+            return _diffengine.make_dense_left_matmul(
+                children[1],
+                A.flatten(order='C'),
+                m,
+                n,
+            )
     elif right_arg.is_constant():
         A = right_arg.value
-
-        if not isinstance(A, sparse.csr_matrix):
-            A = sparse.csr_matrix(A)
-
-        return _diffengine.make_right_matmul(
-            children[0],
-            A.data.astype(np.float64),
-            A.indices.astype(np.int32),
-            A.indptr.astype(np.int32),
-            A.shape[0],
-            A.shape[1],
-        )
+
+        if sparse.issparse(A):
+            if not isinstance(A, sparse.csr_matrix):
+                A = sparse.csr_matrix(A)
+
+            return _diffengine.make_sparse_right_matmul(
+                children[0],
+                A.data.astype(np.float64, copy=False),
+                A.indices.astype(np.int32, copy=False),
+                A.indptr.astype(np.int32, copy=False),
+                A.shape[0],
+                A.shape[1],
+            )
+        else:
+            m, n = normalize_shape(A.shape)
+            return _diffengine.make_dense_right_matmul(
+                children[0],
+                A.flatten(order='C'),
+                m,
+                n,
+            )
     else:
         return _diffengine.make_matmul(children[0], children[1])
 

cvxpy/tests/nlp_tests/derivative_checker.py

@@ -89,7 +89,6 @@ def check_constraint_values(self, x=None):
             python_values.append(constr_val)
 
         python_values = np.hstack(python_values) if python_values else np.array([])
-
         match = np.allclose(c_values, python_values, rtol=1e-10, atol=1e-10)
         return match
 

cvxpy/tests/nlp_tests/stress_tests_diff_engine/test_matmul_sparse.py

@@ -73,3 +73,57 @@ def test_dense_sparse_sparse(self):
         assert np.allclose(dense_val, csc_val)
         assert np.allclose(dense_sol, sparse_sol)
         assert np.allclose(dense_sol, csc_sol)
+
+    def test_dense_left_matmul(self):
+        np.random.seed(0)
+        m, n = 4, 4
+        A = np.random.rand(m, n)
+        X = cp.Variable((n, n), nonneg=True)
+        B = np.random.rand(m, n)
+        obj = cp.Minimize(cp.sum_squares(A @ X - B))
+        constraints = []
+        problem = cp.Problem(obj, constraints)
+        problem.solve(nlp=True, verbose=True)
+        checker = DerivativeChecker(problem)
+        checker.run_and_assert()
+
+    def test_dense_right_matmul(self):
+        np.random.seed(0)
+        m, n = 4, 4
+        A = np.random.rand(m, n)
+        X = cp.Variable((n, n), nonneg=True)
+        B = np.random.rand(m, n)
+        obj = cp.Minimize(cp.sum_squares(X @ A - B))
+        constraints = []
+        problem = cp.Problem(obj, constraints)
+        problem.solve(nlp=True, verbose=True)
+        checker = DerivativeChecker(problem)
+        checker.run_and_assert()
+
+    def test_sparse_and_dense_matmul(self):
+        np.random.seed(0)
+        m, n = 4, 4
+        A = np.random.rand(m, n)
+        C = sp.random(m, n, density=0.5)
+        X = cp.Variable((n, n), nonneg=True)
+        B = np.random.rand(m, n)
+        obj = cp.Minimize(cp.sum_squares(A @ X @ C - B))
+        constraints = []
+        problem = cp.Problem(obj, constraints)
+        problem.solve(nlp=True, verbose=True)
+        checker = DerivativeChecker(problem)
+        checker.run_and_assert()
+
+    def test_sparse_and_dense_matmul2(self):
+        np.random.seed(0)
+        m, n = 4, 3
+        A = np.random.rand(n, m)
+        C = sp.random(m, n, density=0.5)
+        X = cp.Variable((n, n), nonneg=True)
+        B = np.random.rand(m, m)
+        obj = cp.Minimize(cp.sum_squares(C @ X @ A - B))
+        constraints = []
+        problem = cp.Problem(obj, constraints)
+        problem.solve(nlp=True, verbose=True)
+        checker = DerivativeChecker(problem)
+        checker.run_and_assert()