Parameter support

include/affine.h

@@ -32,8 +32,8 @@ expr *new_hstack(expr **args, int n_args, int n_vars);
 expr *new_promote(expr *child, int d1, int d2);
 expr *new_trace(expr *child);
 
-expr *new_constant(int d1, int d2, int n_vars, const double *values);
 expr *new_variable(int d1, int d2, int var_id, int n_vars);
+expr *new_parameter(int d1, int d2, int param_id, int n_vars, const double *values);
 
 expr *new_index(expr *child, int d1, int d2, const int *indices, int n_idxs);
 expr *new_reshape(expr *child, int d1, int d2);

include/bivariate.h

@@ -30,16 +30,16 @@ 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 is a constant matrix */
-expr *new_left_matmul(expr *u, const CSR_Matrix *A);
+/* Left matrix multiplication: A @ f(x) where A comes from a parameter node */
+expr *new_left_matmul(expr *param_node, expr *child, const CSR_Matrix *A);
 
-/* Right matrix multiplication: f(x) @ A where A is a constant matrix */
+/* Right matrix multiplication: f(x) @ A where A is a fixed parameter matrix */
 expr *new_right_matmul(expr *u, const CSR_Matrix *A);
 
-/* Constant scalar multiplication: a * f(x) where a is a constant double */
-expr *new_const_scalar_mult(double a, expr *child);
+/* Scalar multiplication: a * f(x) where a comes from a parameter node */
+expr *new_scalar_mult(expr *param_node, expr *child);
 
-/* Constant vector elementwise multiplication: a ∘ f(x) where a is constant */
-expr *new_const_vector_mult(const double *a, expr *child);
+/* Vector elementwise multiplication: a ∘ f(x) where a comes from a parameter node */
+expr *new_vector_mult(expr *param_node, expr *child);
 
 #endif /* BIVARIATE_H */

include/problem.h

@@ -59,6 +59,11 @@ typedef struct problem
      * hessian are called */
     bool jacobian_called;
 
+    /* Parameter tracking for fast parameter updates */
+    expr **param_nodes; /* weak references to parameter nodes in tree */
+    int n_param_nodes;
+    int total_parameter_size;
+
     /* Statistics for performance measurement */
     Diff_engine_stats stats;
     bool verbose;
@@ -78,4 +83,8 @@ void problem_gradient(problem *prob);
 void problem_jacobian(problem *prob);
 void problem_hessian(problem *prob, double obj_w, const double *w);
 
+/* Parameter support */
+void problem_register_params(problem *prob, expr **param_nodes, int n_param_nodes);
+void problem_update_params(problem *prob, const double *theta);
+
 #endif

include/subexpr.h

@@ -25,6 +25,19 @@
 /* Forward declaration */
 struct int_double_pair;
 
+/* param_id value for fixed (constant) parameters */
+#define PARAM_FIXED -1
+
+/* Parameter node: unified leaf for constants and updatable parameters.
+ * Constants use param_id == PARAM_FIXED and have values set at creation.
+ * Updatable parameters have param_id >= 0 and are updated via problem_update_params.
+ */
+typedef struct parameter_expr
+{
+    expr base;
+    int param_id; /* offset into global theta vector, or PARAM_FIXED */
+} parameter_expr;
+
 /* Type-specific expression structures that "inherit" from expr */
 
 /* Linear operator: y = A * x + b */
@@ -113,6 +126,7 @@ typedef struct left_matmul_expr
     CSC_Matrix *Jchild_CSC;
     CSC_Matrix *J_CSC;
     int *csc_to_csr_workspace;
+    expr *param_source; /* parameter node; A/AT values are refreshed from this */
 } left_matmul_expr;
 
 /* Right matrix multiplication: y = f(x) * A where f(x) is an expression.
@@ -126,19 +140,20 @@ typedef struct right_matmul_expr
     CSC_Matrix *CSC_work;
 } right_matmul_expr;
 
-/* Constant scalar multiplication: y = a * child where a is a constant double */
-typedef struct const_scalar_mult_expr
+/* Scalar multiplication: y = a * child where a comes from a parameter node */
+typedef struct scalar_mult_expr
 {
     expr base;
-    double a;
-} const_scalar_mult_expr;
+    expr *param_source; /* always set; read a from param_source->value[0] */
+} scalar_mult_expr;
 
-/* Constant vector elementwise multiplication: y = a \circ child for constant a */
-typedef struct const_vector_mult_expr
+/* Vector elementwise multiplication: y = a \circ child where a comes from a
+ * parameter node */
+typedef struct vector_mult_expr
 {
     expr base;
-    double *a; /* length equals node->size */
-} const_vector_mult_expr;
+    expr *param_source; /* always set; read a from param_source->value */
+} vector_mult_expr;
 
 /* Index/slicing: y = child[indices] where indices is a list of flat positions */
 typedef struct index_expr

src/affine/hstack.c

@@ -165,13 +165,10 @@ static void free_type_data(expr *node)
     for (int i = 0; i < hnode->n_args; i++)
     {
         free_expr(hnode->args[i]);
-        hnode->args[i] = NULL;
     }
 
     free_csr_matrix(hnode->CSR_work);
-    hnode->CSR_work = NULL;
     free(hnode->args);
-    hnode->args = NULL;
 }
 
 expr *new_hstack(expr **args, int n_args, int n_vars)

src/affine/index.c

@@ -154,11 +154,7 @@ static bool is_affine(const expr *node)
 static void free_type_data(expr *node)
 {
     index_expr *idx = (index_expr *) node;
-    if (idx->indices)
-    {
-        free(idx->indices);
-        idx->indices = NULL;
-    }
+    free(idx->indices);
 }
 
 expr *new_index(expr *child, int d1, int d2, const int *indices, int n_idxs)

src/affine/linear_op.c

@@ -62,11 +62,7 @@ static void free_type_data(expr *node)
     if (lin_node->b != NULL)
     {
         free(lin_node->b);
-        lin_node->b = NULL;
     }
-
-    lin_node->A_csr = NULL;
-    lin_node->A_csc = NULL;
 }
 
 static void jacobian_init(expr *node)

src/affine/constant.c → src/affine/parameter.c

@@ -15,39 +15,48 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
+
+/* Unified parameter/constant leaf node.
+ *
+ * When param_id == PARAM_FIXED, this is a constant whose values are set at
+ * creation and never change. When param_id >= 0, values are updated via
+ * problem_update_params.
+ *
+ * In both cases the derivative behavior is identical: zero Jacobian and
+ * Hessian with respect to variables (always affine). */
+
 #include "affine.h"
+#include "subexpr.h"
 #include <stdlib.h>
 #include <string.h>
 
 static void forward(expr *node, const double *u)
 {
-    /* Constants don't depend on u; values are already set */
+    /* Values are set at creation (constants) or by problem_update_params */
     (void) node;
     (void) u;
 }
 
 static void jacobian_init(expr *node)
 {
-    /* Constant jacobian is all zeros: size x n_vars with 0 nonzeros.
-     * new_csr_matrix uses calloc for row pointers, so they're already 0. */
+    /* Parameter/constant jacobian is all zeros: size x n_vars with 0 nonzeros */
     node->jacobian = new_csr_matrix(node->size, node->n_vars, 0);
 }
 
 static void eval_jacobian(expr *node)
 {
-    /* Constant jacobian never changes - nothing to evaluate */
+    /* Jacobian never changes */
     (void) node;
 }
 
 static void wsum_hess_init(expr *node)
 {
-    /* Constant Hessian is all zeros: n_vars x n_vars with 0 nonzeros. */
+    /* Hessian is all zeros */
     node->wsum_hess = new_csr_matrix(node->n_vars, node->n_vars, 0);
 }
 
 static void eval_wsum_hess(expr *node, const double *w)
 {
-    /* Constant Hessian is always zero - nothing to compute */
     (void) node;
     (void) w;
 }
@@ -58,12 +67,19 @@ static bool is_affine(const expr *node)
     return true;
 }
 
-expr *new_constant(int d1, int d2, int n_vars, const double *values)
+expr *new_parameter(int d1, int d2, int param_id, int n_vars, const double *values)
 {
-    expr *node = (expr *) calloc(1, sizeof(expr));
-    init_expr(node, d1, d2, n_vars, forward, jacobian_init, eval_jacobian, is_affine,
-              wsum_hess_init, eval_wsum_hess, NULL);
-    memcpy(node->value, values, node->size * sizeof(double));
+    parameter_expr *pnode = (parameter_expr *) calloc(1, sizeof(parameter_expr));
+    init_expr(&pnode->base, d1, d2, n_vars, forward, jacobian_init, eval_jacobian,
+              is_affine, wsum_hess_init, eval_wsum_hess, NULL);
+    pnode->param_id = param_id;
+
+    /* If values provided (fixed constant), copy them now */
+    if (values != NULL)
+    {
+        memcpy(pnode->base.value, values, pnode->base.size * sizeof(double));
+    }
+    /* Otherwise values will be populated by problem_update_params */
 
-    return node;
+    return &pnode->base;
 }

src/bivariate/left_matmul.c

@@ -23,10 +23,10 @@
 #include <stdio.h>
 #include <stdlib.h>
 
-/* This file implement the atom 'left_matmul' corresponding to the operation y =
-   A @ f(x), where A is a given matrix and f(x) is an arbitrary expression.
-   Here, f(x) can be a vector-valued expression and a matrix-valued
-   expression. The dimensions are A - m x n, f(x) - n x p, y - m x p.
+/* This file implements the atom 'left_matmul' corresponding to the operation y =
+   A @ f(x), where A is a given matrix (from a parameter node) and f(x) is an
+   arbitrary expression. Here, f(x) can be a vector-valued expression and a
+   matrix-valued expression. The dimensions are A - m x n, f(x) - n x p, y - m x p.
    Note that here A does not have global column indices but it is a local matrix.
    This is an important distinction compared to linear_op_expr.
 
@@ -45,24 +45,38 @@
 
     Working in terms of A_kron unifies the implementation of f(x) being
     vector-valued or matrix-valued.
-
-    I (dance858) think we can get additional big speedups when A is dense by
-    introducing a dense matrix class.
 */
 
 #include "utils/utils.h"
+#include <string.h>
+
+/* Refresh A and AT values from param_source.
+   A is the small m x n matrix (NOT block-diagonal).
+   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;
+
+    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);
+}
 
 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 */
+    refresh_param_values(lin_node);
 
     /* child's forward pass */
     node->left->forward(node->left, u);
 
     /* y = A_kron @ vec(f(x)) */
-    CSR_Matrix *A = ((left_matmul_expr *) node)->A;
-    int n_blocks = ((left_matmul_expr *) node)->n_blocks;
-    block_left_multiply_vec(A, x->value, node->value, n_blocks);
+    block_left_multiply_vec(lin_node->A, x->value, node->value, lin_node->n_blocks);
 }
 
 static bool is_affine(const expr *node)
@@ -78,11 +92,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);
-    lin_node->A = NULL;
-    lin_node->AT = NULL;
-    lin_node->Jchild_CSC = NULL;
-    lin_node->J_CSC = NULL;
-    lin_node->csc_to_csr_workspace = NULL;
+    free_expr(lin_node->param_source);
 }
 
 static void jacobian_init(expr *node)
@@ -109,6 +119,9 @@ 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);
@@ -137,7 +150,7 @@ static void wsum_hess_init(expr *node)
 
 static void eval_wsum_hess(expr *node, const double *w)
 {
-    /* compute A^T w*/
+    /* compute A^T w */
     CSR_Matrix *AT = ((left_matmul_expr *) node)->AT;
     int n_blocks = ((left_matmul_expr *) node)->n_blocks;
     block_left_multiply_vec(AT, w, node->dwork, n_blocks);
@@ -147,51 +160,49 @@ static void eval_wsum_hess(expr *node, const double *w)
            node->wsum_hess->nnz * sizeof(double));
 }
 
-expr *new_left_matmul(expr *u, const CSR_Matrix *A)
+expr *new_left_matmul(expr *param_node, expr *child, const CSR_Matrix *A)
 {
-    /* We expect u->d1 == A->n. However, numpy's broadcasting rules allow users
-       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 A_m = A->m;
+    int A_n = A->n;
+
+    /* Dimension logic: handle numpy broadcasting (1, n) as (n, ) */
     int d1, d2, n_blocks;
-    if (u->d1 == A->n)
+    if (child->d1 == A_n)
     {
-        d1 = A->m;
-        d2 = u->d2;
-        n_blocks = u->d2;
+        d1 = A_m;
+        d2 = child->d2;
+        n_blocks = child->d2;
     }
-    else if (u->d2 == A->n && u->d1 == 1)
+    else if (child->d2 == A_n && child->d1 == 1)
     {
         d1 = 1;
-        d2 = A->m;
+        d2 = A_m;
         n_blocks = 1;
     }
     else
     {
-        fprintf(stderr, "Error in new_left_matmul: dimension mismatch \n");
+        fprintf(stderr, "Error in new_left_matmul: dimension mismatch\n");
         exit(1);
     }
 
     /* Allocate the type-specific struct */
     left_matmul_expr *lin_node =
         (left_matmul_expr *) calloc(1, sizeof(left_matmul_expr));
     expr *node = &lin_node->base;
-    init_expr(node, d1, d2, u->n_vars, forward, jacobian_init, eval_jacobian,
+    init_expr(node, d1, d2, child->n_vars, forward, jacobian_init, eval_jacobian,
               is_affine, wsum_hess_init, eval_wsum_hess, free_type_data);
-    node->left = u;
-    expr_retain(u);
-
-    /* allocate workspace. iwork is used for transposing A (requiring size A->n)
-       and for converting J_child csr to csc (requring size node->n_vars).
-       csc_to_csr_workspace is used for converting J_CSC to CSR (requring node->size)
-     */
-    node->iwork = (int *) malloc(MAX(A->n, node->n_vars) * sizeof(int));
+    node->left = child;
+    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));
     lin_node->csc_to_csr_workspace = (int *) malloc(node->size * sizeof(int));
     lin_node->n_blocks = n_blocks;
-
-    /* store A and AT */
     lin_node->A = new_csr(A);
     lin_node->AT = transpose(lin_node->A, node->iwork);
 
+    lin_node->param_source = param_node;
+    if (param_node) expr_retain(param_node);
+
     return node;
 }