fix AT workspace in left_matmul and add has_been_refreshed

Author: dance858

include/bivariate.h

@@ -30,7 +30,8 @@ expr *new_rel_entr_second_arg_scalar(expr *left, expr *right);
 /* Matrix multiplication: Z = X @ Y */
 expr *new_matmul(expr *x, expr *y);
 
-/* Left matrix multiplication: A @ f(x) where A comes from a parameter node */
+/* Left matrix multiplication: A @ f(x) where A comes from a parameter node.
+   Only the forward pass possibly updates the parameter. */
 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 */

include/problem.h

@@ -59,7 +59,7 @@ typedef struct problem
      * hessian are called */
     bool jacobian_called;
 
-    /* Parameter tracking for fast parameter updates */
+    /* Parameter tracking for fast parameter updates. */
     expr **param_nodes; /* weak references to parameter nodes in tree */
     int n_param_nodes;
     int total_parameter_size;

include/subexpr.h

@@ -35,7 +35,8 @@ struct int_double_pair;
 typedef struct parameter_expr
 {
     expr base;
-    int param_id; /* offset into global theta vector, or PARAM_FIXED */
+    int param_id;            /* offset into global theta vector, or PARAM_FIXED */
+    bool has_been_refreshed; /* tracks whether parameter has been refreshed */
 } parameter_expr;
 
 /* Type-specific expression structures that "inherit" from expr */
@@ -126,6 +127,7 @@ typedef struct left_matmul_expr
     CSC_Matrix *Jchild_CSC;
     CSC_Matrix *J_CSC;
     int *csc_to_csr_workspace;
+    int *AT_iwork;      /* work for computing AT values from A */
     expr *param_source; /* parameter node; A/AT values are refreshed from this */
 } left_matmul_expr;
 

src/affine/parameter.c

@@ -74,6 +74,7 @@ expr *new_parameter(int d1, int d2, int param_id, int n_vars, const double *valu
     init_expr(node, d1, d2, n_vars, forward, jacobian_init, eval_jacobian, is_affine,
               wsum_hess_init, eval_wsum_hess, NULL);
     pnode->param_id = param_id;
+    pnode->has_been_refreshed = false;
 
     /* If values provided (fixed constant), copy them now.
        Otherwise values will be populated by problem_update_params. */

src/bivariate/left_matmul.c

@@ -55,21 +55,25 @@
    No-op when param_source is NULL (fixed constant — values already in A). */
 static void refresh_param_values(left_matmul_expr *lin_node)
 {
-    if (!lin_node->param_source) return;
+    parameter_expr *param = (parameter_expr *) lin_node->param_source;
 
+    if (!param || param->has_been_refreshed) return;
+    param->has_been_refreshed = true;
+
+    /* update values of A */
     memcpy(lin_node->A->x, lin_node->param_source->value,
            lin_node->A->nnz * sizeof(double));
 
-    /* Recompute AT values from updated A */
-    AT_fill_values(lin_node->A, lin_node->AT, lin_node->base.iwork);
+    /* update values of AT */
+    AT_fill_values(lin_node->A, lin_node->AT, lin_node->AT_iwork);
 }
 
 static void forward(expr *node, const double *u)
 {
     expr *x = node->left;
     left_matmul_expr *lin_node = (left_matmul_expr *) node;
 
-    /* refresh A/AT from parameter source */
+    /* possibly refresh A and AT */
     refresh_param_values(lin_node);
 
     /* child's forward pass */
@@ -92,6 +96,7 @@ static void free_type_data(expr *node)
     free_csc_matrix(lin_node->Jchild_CSC);
     free_csc_matrix(lin_node->J_CSC);
     free(lin_node->csc_to_csr_workspace);
+    free(lin_node->AT_iwork);
     free_expr(lin_node->param_source);
 }
 
@@ -119,9 +124,6 @@ static void eval_jacobian(expr *node)
     CSC_Matrix *Jchild_CSC = lnode->Jchild_CSC;
     CSC_Matrix *J_CSC = lnode->J_CSC;
 
-    /* refresh A from parameter source */
-    refresh_param_values(lnode);
-
     /* evaluate child's jacobian and convert to CSC */
     x->eval_jacobian(x);
     csr_to_csc_fill_values(x->jacobian, Jchild_CSC, node->iwork);
@@ -167,7 +169,6 @@ expr *new_left_matmul(expr *param_node, expr *child, const CSR_Matrix *A)
        to do A @ u where u is (n, ) which in C is actually (1, n). In that case
        the result of A @ u is (m, ), which is (1, m) according to broadcasting
        rules. We therefore check if this is the case. */
-
     int d1, d2, n_blocks;
     if (child->d1 == A->n)
     {
@@ -197,13 +198,14 @@ expr *new_left_matmul(expr *param_node, expr *child, const CSR_Matrix *A)
     expr_retain(child);
 
     /* Store small A (NOT block-diagonal) — block functions handle the rest */
-    node->iwork = (int *) malloc(MAX(A->n, node->n_vars) * sizeof(int));
+    node->iwork = (int *) malloc(node->n_vars * sizeof(int));
+    lin_node->AT_iwork = (int *) malloc(A->n * sizeof(int));
     lin_node->csc_to_csr_workspace = (int *) malloc(node->size * sizeof(int));
     lin_node->n_blocks = n_blocks;
     lin_node->A = new_csr(A);
-    lin_node->AT = transpose(lin_node->A, node->iwork);
-
+    lin_node->AT = transpose(lin_node->A, lin_node->AT_iwork);
     lin_node->param_source = param_node;
+
     if (param_node) expr_retain(param_node);
 
     return node;

src/problem.c

@@ -456,7 +456,9 @@ void problem_register_params(problem *prob, expr **param_nodes, int n_param_node
 
     prob->total_parameter_size = 0;
     for (int i = 0; i < n_param_nodes; i++)
+    {
         prob->total_parameter_size += param_nodes[i]->size;
+    }
 }
 
 void problem_update_params(problem *prob, const double *theta)
@@ -466,7 +468,9 @@ void problem_update_params(problem *prob, const double *theta)
         parameter_expr *p = (parameter_expr *) prob->param_nodes[i];
         if (p->param_id == PARAM_FIXED) continue;
         memcpy(p->base.value, theta + p->param_id, p->base.size * sizeof(double));
+        p->has_been_refreshed = false;
     }
-    /* Force re-evaluation of affine Jacobians on next call */
+
+    /* force re-evaluation of affine Jacobians on next call */
     prob->jacobian_called = false;
 }