Simplify new_left_matmul: accept CSR directly, remove sparse/dense branching

Change signature from (param_node, child) to (param_node, child, A) where
A is a const CSR_Matrix providing the sparsity pattern and initial values.
The constructor copies A with new_csr() instead of rebuilding CSR from
dense column-major values. Remove src_m/src_n fields from left_matmul_expr
(read dimensions from A->m instead). Update all call sites and tests.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Author: Transurgeon

include/bivariate.h

@@ -31,7 +31,7 @@ expr *new_rel_entr_second_arg_scalar(expr *left, expr *right);
 expr *new_matmul(expr *x, expr *y);
 
 /* Left matrix multiplication: A @ f(x) where A comes from a parameter node */
-expr *new_left_matmul(expr *param_node, expr *child);
+expr *new_left_matmul(expr *param_node, expr *child, const CSR_Matrix *A);
 
 /* Right matrix multiplication: f(x) @ A where A is a fixed parameter matrix */
 expr *new_right_matmul(expr *u, const CSR_Matrix *A);

include/subexpr.h

@@ -127,7 +127,6 @@ typedef struct left_matmul_expr
     CSC_Matrix *J_CSC;
     int *csc_to_csr_workspace;
     expr *param_source; /* parameter node; A/AT values are refreshed from this */
-    int src_m, src_n;   /* original matrix dimensions */
 } left_matmul_expr;
 
 /* Right matrix multiplication: y = f(x) * A where f(x) is an expression.

src/bivariate/left_matmul.c

@@ -55,7 +55,7 @@
 static void refresh_param_values(left_matmul_expr *lin_node)
 {
     const double *src = lin_node->param_source->value;
-    int m = lin_node->src_m;
+    int m = lin_node->A->m;
     CSR_Matrix *A = lin_node->A;
 
     /* Fill A values from column-major source, following existing sparsity pattern */
@@ -163,10 +163,10 @@ static void eval_wsum_hess(expr *node, const double *w)
            node->wsum_hess->nnz * sizeof(double));
 }
 
-expr *new_left_matmul(expr *param_node, expr *child)
+expr *new_left_matmul(expr *param_node, expr *child, const CSR_Matrix *A)
 {
-    int A_m = param_node->d1;
-    int A_n = param_node->d2;
+    int A_m = A->m;
+    int A_n = A->n;
 
     /* Dimension logic: handle numpy broadcasting (1, n) as (n, ) */
     int d1, d2, n_blocks;
@@ -188,40 +188,6 @@ expr *new_left_matmul(expr *param_node, expr *child)
         exit(1);
     }
 
-    /* Build CSR from param_node's column-major values.
-     * For fixed parameters (PARAM_FIXED), skip zeros to preserve sparsity.
-     * For updatable parameters, build dense CSR since sparsity may change. */
-    parameter_expr *pnode = (parameter_expr *) param_node;
-    int sparse = (pnode->param_id == PARAM_FIXED);
-
-    int nnz = 0;
-    if (sparse)
-    {
-        for (int row = 0; row < A_m; row++)
-            for (int col = 0; col < A_n; col++)
-                if (param_node->value[row + col * A_m] != 0.0) nnz++;
-    }
-    else
-    {
-        nnz = A_m * A_n;
-    }
-
-    CSR_Matrix *A = new_csr_matrix(A_m, A_n, nnz);
-    int idx = 0;
-    for (int row = 0; row < A_m; row++)
-    {
-        A->p[row] = idx;
-        for (int col = 0; col < A_n; col++)
-        {
-            double val = param_node->value[row + col * A_m];
-            if (sparse && val == 0.0) continue;
-            A->i[idx] = col;
-            A->x[idx] = val;
-            idx++;
-        }
-    }
-    A->p[A_m] = idx;
-
     /* Allocate the type-specific struct */
     left_matmul_expr *lin_node =
         (left_matmul_expr *) calloc(1, sizeof(left_matmul_expr));
@@ -235,12 +201,10 @@ expr *new_left_matmul(expr *param_node, expr *child)
     node->iwork = (int *) malloc(MAX(A_n, node->n_vars) * sizeof(int));
     lin_node->csc_to_csr_workspace = (int *) malloc(node->size * sizeof(int));
     lin_node->n_blocks = n_blocks;
-    lin_node->A = A; /* transfer ownership */
-    lin_node->AT = transpose(A, node->iwork);
+    lin_node->A = new_csr(A);
+    lin_node->AT = transpose(lin_node->A, node->iwork);
 
     lin_node->param_source = param_node;
-    lin_node->src_m = A_m;
-    lin_node->src_n = A_n;
     expr_retain(param_node);
 
     return node;

src/bivariate/right_matmul.c

@@ -43,7 +43,7 @@ expr *new_right_matmul(expr *u, const CSR_Matrix *A)
 
     expr *u_transpose = new_transpose(u);
     expr *param_node = new_parameter(m, n, PARAM_FIXED, u->n_vars, col_major);
-    expr *left_matmul_node = new_left_matmul(param_node, u_transpose);
+    expr *left_matmul_node = new_left_matmul(param_node, u_transpose, AT);
     expr *node = new_transpose(left_matmul_node);
 
     free(col_major);

tests/jacobian_tests/test_left_matmul.h

@@ -1,5 +1,6 @@
 #include <math.h>
 #include <stdio.h>
+#include <string.h>
 
 #include "bivariate.h"
 #include "elementwise_univariate.h"
@@ -36,7 +37,16 @@ const char *test_jacobian_left_matmul_log()
     expr *A_param = new_parameter(4, 3, PARAM_FIXED, 3, A_vals);
 
     expr *log_x = new_log(x);
-    expr *A_log_x = new_left_matmul(A_param, log_x);
+    CSR_Matrix *A_csr = new_csr_matrix(4, 3, 7);
+    int A_p[5] = {0, 2, 4, 6, 7};
+    int A_i[7] = {0, 2, 0, 2, 0, 2, 0};
+    double A_x[7] = {1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0};
+    memcpy(A_csr->p, A_p, 5 * sizeof(int));
+    memcpy(A_csr->i, A_i, 7 * sizeof(int));
+    memcpy(A_csr->x, A_x, 7 * sizeof(double));
+
+    expr *A_log_x = new_left_matmul(A_param, log_x, A_csr);
+    free_csr_matrix(A_csr);
 
     A_log_x->forward(A_log_x, x_vals);
     A_log_x->jacobian_init(A_log_x);
@@ -74,7 +84,16 @@ const char *test_jacobian_left_matmul_log_matrix()
     expr *A_param = new_parameter(4, 3, PARAM_FIXED, 6, A_vals);
 
     expr *log_x = new_log(x);
-    expr *A_log_x = new_left_matmul(A_param, log_x);
+    CSR_Matrix *A_csr = new_csr_matrix(4, 3, 7);
+    int A_p[5] = {0, 2, 4, 6, 7};
+    int A_i[7] = {0, 2, 0, 2, 0, 2, 0};
+    double A_x[7] = {1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0};
+    memcpy(A_csr->p, A_p, 5 * sizeof(int));
+    memcpy(A_csr->i, A_i, 7 * sizeof(int));
+    memcpy(A_csr->x, A_x, 7 * sizeof(double));
+
+    expr *A_log_x = new_left_matmul(A_param, log_x, A_csr);
+    free_csr_matrix(A_csr);
 
     A_log_x->forward(A_log_x, x_vals);
     A_log_x->jacobian_init(A_log_x);
@@ -140,7 +159,16 @@ const char *test_jacobian_left_matmul_log_composite()
 
     expr *Bx = new_linear(x, B, NULL);
     expr *log_Bx = new_log(Bx);
-    expr *A_log_Bx = new_left_matmul(A_param, log_Bx);
+    CSR_Matrix *A_csr = new_csr_matrix(4, 3, 7);
+    int A_p[5] = {0, 2, 4, 6, 7};
+    int A_i[7] = {0, 2, 0, 2, 0, 2, 0};
+    double A_x[7] = {1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0};
+    memcpy(A_csr->p, A_p, 5 * sizeof(int));
+    memcpy(A_csr->i, A_i, 7 * sizeof(int));
+    memcpy(A_csr->x, A_x, 7 * sizeof(double));
+
+    expr *A_log_Bx = new_left_matmul(A_param, log_Bx, A_csr);
+    free_csr_matrix(A_csr);
 
     A_log_Bx->forward(A_log_Bx, x_vals);
     A_log_Bx->jacobian_init(A_log_Bx);

tests/jacobian_tests/test_transpose.h

@@ -3,11 +3,13 @@
 #define TEST_TRANSPOSE_H
 
 #include "affine.h"
+#include "bivariate.h"
 #include "minunit.h"
 #include "subexpr.h"
 #include "test_helpers.h"
 #include <math.h>
 #include <stdio.h>
+#include <string.h>
 
 const char *test_jacobian_transpose()
 {
@@ -17,7 +19,17 @@ const char *test_jacobian_transpose()
 
     // X = [1 2; 3 4] (columnwise: x = [1 3 2 4])
     expr *X = new_variable(2, 2, 0, 4);
-    expr *AX = new_left_matmul(A_param, X);
+    /* Build dense 2x2 CSR */
+    CSR_Matrix *A_csr = new_csr_matrix(2, 2, 4);
+    int Ap[3] = {0, 2, 4};
+    int Ai[4] = {0, 1, 0, 1};
+    double Ax[4] = {1.0, 2.0, 3.0, 4.0};
+    memcpy(A_csr->p, Ap, 3 * sizeof(int));
+    memcpy(A_csr->i, Ai, 4 * sizeof(int));
+    memcpy(A_csr->x, Ax, 4 * sizeof(double));
+
+    expr *AX = new_left_matmul(A_param, X, A_csr);
+    free_csr_matrix(A_csr);
     expr *transpose_AX = new_transpose(AX);
     double u[4] = {1, 3, 2, 4};
     transpose_AX->forward(transpose_AX, u);

tests/problem/test_param_prob.h

@@ -3,6 +3,7 @@
 
 #include <math.h>
 #include <stdio.h>
+#include <string.h>
 
 #include "affine.h"
 #include "bivariate.h"
@@ -185,7 +186,17 @@ const char *test_param_left_matmul_problem(void)
     /* Constraint: A @ x */
     expr *x_con = new_variable(2, 1, 0, n_vars);
     expr *A_param = new_parameter(2, 2, 0, n_vars, NULL);
-    expr *constraint = new_left_matmul(A_param, x_con);
+
+    CSR_Matrix *A_csr = new_csr_matrix(2, 2, 4);
+    int Ap[3] = {0, 2, 4};
+    int Ai[4] = {0, 1, 0, 1};
+    double Ax[4] = {0.0, 0.0, 0.0, 0.0};
+    memcpy(A_csr->p, Ap, 3 * sizeof(int));
+    memcpy(A_csr->i, Ai, 4 * sizeof(int));
+    memcpy(A_csr->x, Ax, 4 * sizeof(double));
+
+    expr *constraint = new_left_matmul(A_param, x_con, A_csr);
+    free_csr_matrix(A_csr);
 
     expr *constraints[1] = {constraint};
 

tests/profiling/profile_left_matmul.h

@@ -28,7 +28,26 @@ const char *profile_left_matmul()
     expr *A_param = new_parameter(n, n, PARAM_FIXED, n, A_vals);
     free(A_vals);
 
-    expr *AX = new_left_matmul(A_param, X);
+    /* Build dense n x n CSR (all ones) */
+    int nnz = n * n;
+    CSR_Matrix *A_csr = new_csr_matrix(n, n, nnz);
+    {
+        int idx = 0;
+        for (int row = 0; row < n; row++)
+        {
+            A_csr->p[row] = idx;
+            for (int col = 0; col < n; col++)
+            {
+                A_csr->i[idx] = col;
+                A_csr->x[idx] = 1.0;
+                idx++;
+            }
+        }
+        A_csr->p[n] = idx;
+    }
+
+    expr *AX = new_left_matmul(A_param, X, A_csr);
+    free_csr_matrix(A_csr);
 
     double *x_vals = (double *) malloc(n * n * sizeof(double));
     for (int i = 0; i < n * n; i++)

tests/wsum_hess/test_left_matmul.h

@@ -58,7 +58,16 @@ const char *test_wsum_hess_left_matmul()
     expr *A_param = new_parameter(4, 3, PARAM_FIXED, 3, A_vals);
 
     expr *log_x = new_log(x);
-    expr *A_log_x = new_left_matmul(A_param, log_x);
+    CSR_Matrix *A_csr = new_csr_matrix(4, 3, 7);
+    int A_p[5] = {0, 2, 4, 6, 7};
+    int A_i[7] = {0, 2, 0, 2, 0, 2, 0};
+    double A_x[7] = {1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0};
+    memcpy(A_csr->p, A_p, 5 * sizeof(int));
+    memcpy(A_csr->i, A_i, 7 * sizeof(int));
+    memcpy(A_csr->x, A_x, 7 * sizeof(double));
+
+    expr *A_log_x = new_left_matmul(A_param, log_x, A_csr);
+    free_csr_matrix(A_csr);
 
     A_log_x->forward(A_log_x, x_vals);
     A_log_x->jacobian_init(A_log_x);
@@ -151,7 +160,16 @@ const char *test_wsum_hess_left_matmul_composite()
 
     expr *Bx = new_linear(x, B, NULL);
     expr *log_Bx = new_log(Bx);
-    expr *A_log_Bx = new_left_matmul(A_param, log_Bx);
+    CSR_Matrix *A_csr = new_csr_matrix(4, 3, 7);
+    int A_p[5] = {0, 2, 4, 6, 7};
+    int A_i[7] = {0, 2, 0, 2, 0, 2, 0};
+    double A_x[7] = {1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0};
+    memcpy(A_csr->p, A_p, 5 * sizeof(int));
+    memcpy(A_csr->i, A_i, 7 * sizeof(int));
+    memcpy(A_csr->x, A_x, 7 * sizeof(double));
+
+    expr *A_log_Bx = new_left_matmul(A_param, log_Bx, A_csr);
+    free_csr_matrix(A_csr);
 
     A_log_Bx->forward(A_log_Bx, x_vals);
     A_log_Bx->jacobian_init(A_log_Bx);
@@ -217,7 +235,16 @@ const char *test_wsum_hess_left_matmul_matrix()
     expr *A_param = new_parameter(4, 3, PARAM_FIXED, 6, A_vals);
 
     expr *log_x = new_log(x);
-    expr *A_log_x = new_left_matmul(A_param, log_x);
+    CSR_Matrix *A_csr = new_csr_matrix(4, 3, 7);
+    int A_p[5] = {0, 2, 4, 6, 7};
+    int A_i[7] = {0, 2, 0, 2, 0, 2, 0};
+    double A_x[7] = {1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0};
+    memcpy(A_csr->p, A_p, 5 * sizeof(int));
+    memcpy(A_csr->i, A_i, 7 * sizeof(int));
+    memcpy(A_csr->x, A_x, 7 * sizeof(double));
+
+    expr *A_log_x = new_left_matmul(A_param, log_x, A_csr);
+    free_csr_matrix(A_csr);
 
     A_log_x->forward(A_log_x, x_vals);
     A_log_x->jacobian_init(A_log_x);