new parameter abstraction that fixess the fialing test

Author: dance858

include/expr.h

@@ -86,6 +86,12 @@ typedef struct expr
     local_wsum_hess_fn local_wsum_hess; /* used by elementwise univariate atoms*/
     free_type_data_fn free_type_data;   /* Cleanup for type-specific fields */
     Expr_Work *work;                    /* derivative workspace */
+    /* Set to true on all nodes by problem_update_params() via
+       expr_set_needs_refresh(). Atoms that cache parameter data
+       (e.g. left_matmul_dense) check this flag before their forward
+       pass: if true, they refresh their cached matrices from
+       param_source->value and clear the flag to false. */
+    bool needs_parameter_refresh;
 
     // name of node just for debugging - should be removed later
     char name[32];
@@ -108,6 +114,9 @@ void wsum_hess_init(expr *node);
  * Must be called after jacobian_init. */
 void jacobian_csc_init(expr *node);
 
+/* Recursively set needs_parameter_refresh on node and all children */
+void expr_set_needs_refresh(expr *node);
+
 /* Reference counting helpers */
 void expr_retain(expr *node);
 

include/subexpr.h

@@ -37,9 +37,6 @@ typedef struct parameter_expr
 {
     expr base;
     int param_id;
-    /* Set to true by problem_update_params(), cleared by
-       refresh_param_values() after propagating new values. */
-    bool needs_refresh;
 } parameter_expr;
 
 /* Linear operator: y = A * x + b

src/atoms/affine/left_matmul.c

@@ -55,12 +55,11 @@ static void refresh_param_values(left_matmul_expr *lnode)
     {
         return;
     }
-    parameter_expr *param = (parameter_expr *) lnode->param_source;
-    if (!param->needs_refresh)
+    if (!lnode->base.needs_parameter_refresh)
     {
         return;
     }
-    param->needs_refresh = false;
+    lnode->base.needs_parameter_refresh = false;
     lnode->refresh_param_values(lnode);
 }
 

src/atoms/affine/parameter.c

@@ -65,7 +65,6 @@ expr *new_parameter(int d1, int d2, int param_id, int n_vars, const double *valu
               is_affine, wsum_hess_init_impl, eval_wsum_hess, NULL);
 
     pnode->param_id = param_id;
-    pnode->needs_refresh = false;
 
     if (values != NULL)
     {

src/expr.c

@@ -108,6 +108,14 @@ void wsum_hess_init(expr *node)
     node->wsum_hess_init_impl(node);
 }
 
+void expr_set_needs_refresh(expr *node)
+{
+    if (node == NULL) return;
+    node->needs_parameter_refresh = true;
+    expr_set_needs_refresh(node->left);
+    expr_set_needs_refresh(node->right);
+}
+
 void expr_retain(expr *node)
 {
     if (node == NULL) return;

src/problem.c

@@ -382,7 +382,13 @@ void problem_update_params(problem *prob, const double *theta)
         if (param->param_id == PARAM_FIXED) continue;
         int offset = param->param_id;
         memcpy(pnode->value, theta + offset, pnode->size * sizeof(double));
-        param->needs_refresh = true;
+    }
+
+    /* Propagate needs_parameter_refresh to all expressions */
+    expr_set_needs_refresh(prob->objective);
+    for (int i = 0; i < prob->n_constraints; i++)
+    {
+        expr_set_needs_refresh(prob->constraints[i]);
     }
 
     /* Force re-evaluation of affine Jacobians on next call */