Merge origin/main into parameter-support-v2

Resolves conflicts from main's workspace refactor (iwork/dwork -> work->),
elementwise atom split (elementwise_univariate -> full_dom/restricted_dom),
chain rule refactor, vstack addition, and numerical diff checker.

Removed stale test_jacobian_composite_log tests (incompatible with main's
chain rule refactor; covered by test_chain_rule_jacobian and numerical diff).

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Author: Transurgeon

CMakeLists.txt

@@ -103,6 +103,7 @@ if(NOT SKBUILD)
     add_executable(all_tests
         tests/all_tests.c
         tests/test_helpers.c
+        tests/numerical_diff.c
     )
     target_link_libraries(all_tests dnlp_diff)
 

include/affine.h

@@ -29,6 +29,7 @@ expr *new_neg(expr *child);
 
 expr *new_sum(expr *child, int axis);
 expr *new_hstack(expr **args, int n_args, int n_vars);
+expr *new_vstack(expr **args, int n_args, int n_vars);
 expr *new_promote(expr *child, int d1, int d2);
 expr *new_trace(expr *child);
 

include/elementwise_full_dom.h

@@ -0,0 +1,34 @@
+#ifndef ELEMENTWISE_FULL_DOM_H
+#define ELEMENTWISE_FULL_DOM_H
+
+#include "expr.h"
+
+/* Helper function to initialize an elementwise expr
+ * (can be used with derived types) */
+void init_elementwise(expr *node, expr *child);
+
+expr *new_exp(expr *child);
+expr *new_sin(expr *child);
+expr *new_cos(expr *child);
+expr *new_sinh(expr *child);
+expr *new_tanh(expr *child);
+expr *new_asinh(expr *child);
+expr *new_logistic(expr *child);
+expr *new_power(expr *child, double p);
+expr *new_xexp(expr *child);
+expr *new_normal_cdf(expr *child);
+
+/* the jacobian and wsum_hess for elementwise full domain
+   atoms are always initialized in the same way and
+   implement the chain rule in the same way */
+void jacobian_init_elementwise(expr *node);
+void eval_jacobian_elementwise(expr *node);
+void wsum_hess_init_elementwise(expr *node);
+void eval_wsum_hess_elementwise(expr *node, const double *w);
+expr *new_elementwise(expr *child);
+
+/* no elementwise atoms are affine according to our
+   convention, so we can have a common implementation */
+bool is_affine_elementwise(const expr *node);
+
+#endif /* ELEMENTWISE_FULL_DOM_H */

include/elementwise_restricted_dom.h

@@ -0,0 +1,18 @@
+#ifndef ELEMENTWISE_RESTRICTED_DOM_H
+#define ELEMENTWISE_RESTRICTED_DOM_H
+
+#include "expr.h"
+
+/* Shared init functions for restricted domain atoms
+ * (variable-child only, no linear operator support) */
+void jacobian_init_restricted(expr *node);
+void wsum_hess_init_restricted(expr *node);
+bool is_affine_restricted(const expr *node);
+expr *new_restricted(expr *child);
+
+expr *new_log(expr *child);
+expr *new_entr(expr *child);
+expr *new_atanh(expr *child);
+expr *new_tan(expr *child);
+
+#endif /* ELEMENTWISE_RESTRICTED_DOM_H */

include/elementwise_univariate.h

@@ -39,7 +39,24 @@ typedef void (*local_wsum_hess_fn)(struct expr *node, double *out, const double
 typedef bool (*is_affine_fn)(const struct expr *node);
 typedef void (*free_type_data_fn)(struct expr *node);
 
-/* Base expression node structure - contains only common fields */
+/* Workspace for derivative computation */
+typedef struct
+{
+    double *dwork;
+    int *iwork;
+    CSC_Matrix *jacobian_csc;
+    int *csc_work; /* for CSR-CSC conversion */
+
+    /* jacobian_csc_filled is only used for affine functions to avoid redundant
+       conversions. Could become relevant for non-affine functions if we start
+       supporting common subexpressions on the Python side. */
+    bool jacobian_csc_filled;
+    double *local_jac_diag; /* cached f'(g(x)) diagonal */
+    CSR_Matrix *hess_term1; /* Jg^T D Jg workspace */
+    CSR_Matrix *hess_term2; /* child wsum_hess workspace */
+} Expr_Work;
+
+/* Base expression node structure */
 typedef struct expr
 {
     // ------------------------------------------------------------------------
@@ -48,8 +65,6 @@ typedef struct expr
     int d1, d2, size, n_vars, refcount, var_id;
     struct expr *left;
     struct expr *right;
-    double *dwork;
-    int *iwork;
 
     // ------------------------------------------------------------------------
     //                     oracle related quantities
@@ -70,6 +85,7 @@ typedef struct expr
     local_jacobian_fn local_jacobian;   /* used by elementwise univariate atoms*/
     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 */
 
     // name of node just for debugging - should be removed later
     char name[32];
@@ -83,6 +99,10 @@ void init_expr(expr *node, int d1, int d2, int n_vars, forward_fn forward,
 
 void free_expr(expr *node);
 
+/* Initialize CSC form of the Jacobian from the CSR Jacobian.
+ * Must be called after jacobian_init. */
+void jacobian_csc_init(expr *node);
+
 /* Reference counting helpers */
 void expr_retain(expr *node);
 

include/expr.h

@@ -28,6 +28,7 @@ typedef struct CSR_Matrix
 /* constructors and destructors */
 CSR_Matrix *new_csr_matrix(int m, int n, int nnz);
 CSR_Matrix *new_csr(const CSR_Matrix *A);
+CSR_Matrix *new_csr_copy_sparsity(const CSR_Matrix *A);
 void free_csr_matrix(CSR_Matrix *matrix);
 void copy_csr_matrix(const CSR_Matrix *A, CSR_Matrix *C);
 

include/utils/CSR_Matrix.h

@@ -191,12 +191,10 @@ static void wsum_hess_init(expr *node)
     x->wsum_hess_init(x);
 
     /* Same sparsity as child - weights get summed */
-    node->wsum_hess = new_csr_matrix(node->n_vars, node->n_vars, x->wsum_hess->nnz);
-    memcpy(node->wsum_hess->p, x->wsum_hess->p, (x->wsum_hess->m + 1) * sizeof(int));
-    memcpy(node->wsum_hess->i, x->wsum_hess->i, x->wsum_hess->nnz * sizeof(int));
+    node->wsum_hess = new_csr_copy_sparsity(x->wsum_hess);
 
     /* allocate space for weight vector */
-    node->dwork = malloc(node->size * sizeof(double));
+    node->work->dwork = malloc(node->size * sizeof(double));
 }
 
 static void eval_wsum_hess(expr *node, const double *w)
@@ -205,7 +203,7 @@ static void eval_wsum_hess(expr *node, const double *w)
     expr *x = node->left;
 
     /* Zero out the work array first */
-    memset(node->dwork, 0, x->size * sizeof(double));
+    memset(node->work->dwork, 0, x->size * sizeof(double));
 
     if (bcast->type == BROADCAST_ROW)
     {
@@ -214,7 +212,7 @@ static void eval_wsum_hess(expr *node, const double *w)
         {
             for (int i = 0; i < node->d1; i++)
             {
-                node->dwork[j] += w[i + j * node->d1];
+                node->work->dwork[j] += w[i + j * node->d1];
             }
         }
     }
@@ -225,21 +223,21 @@ static void eval_wsum_hess(expr *node, const double *w)
         {
             for (int i = 0; i < node->d1; i++)
             {
-                node->dwork[i] += w[i + j * node->d1];
+                node->work->dwork[i] += w[i + j * node->d1];
             }
         }
     }
     else
     {
         /* (1, 1) -> (m, n): scalar has m*n weights to sum */
-        node->dwork[0] = 0.0;
+        node->work->dwork[0] = 0.0;
         for (int k = 0; k < node->size; k++)
         {
-            node->dwork[0] += w[k];
+            node->work->dwork[0] += w[k];
         }
     }
 
-    x->eval_wsum_hess(x, node->dwork);
+    x->eval_wsum_hess(x, node->work->dwork);
     memcpy(node->wsum_hess->x, x->wsum_hess->x, x->wsum_hess->nnz * sizeof(double));
 }
 

src/affine/broadcast.c

@@ -100,14 +100,12 @@ static void wsum_hess_init(expr *node)
     x->wsum_hess_init(x);
 
     /* workspace for extracting diagonal weights */
-    node->dwork = (double *) calloc(x->size, sizeof(double));
+    node->work->dwork = (double *) calloc(x->size, sizeof(double));
 
     /* Copy child's Hessian structure (diag_vec is linear, so its own Hessian is
      * zero) */
     CSR_Matrix *Hx = x->wsum_hess;
-    node->wsum_hess = new_csr_matrix(Hx->m, Hx->n, Hx->nnz);
-    memcpy(node->wsum_hess->p, Hx->p, (Hx->m + 1) * sizeof(int));
-    memcpy(node->wsum_hess->i, Hx->i, Hx->nnz * sizeof(int));
+    node->wsum_hess = new_csr_copy_sparsity(Hx);
 }
 
 static void eval_wsum_hess(expr *node, const double *w)
@@ -118,11 +116,11 @@ static void eval_wsum_hess(expr *node, const double *w)
     /* Extract weights from diagonal positions of w (which has n^2 elements) */
     for (int i = 0; i < n; i++)
     {
-        node->dwork[i] = w[i * (n + 1)];
+        node->work->dwork[i] = w[i * (n + 1)];
     }
 
     /* Evaluate child's Hessian with extracted weights */
-    x->eval_wsum_hess(x, node->dwork);
+    x->eval_wsum_hess(x, node->work->dwork);
     memcpy(node->wsum_hess->x, x->wsum_hess->x, x->wsum_hess->nnz * sizeof(double));
 }
 

src/affine/diag_vec.c

@@ -104,7 +104,7 @@ static void wsum_hess_init(expr *node)
     x->wsum_hess_init(x);
 
     /* for setting weight vector to evaluate hessian of child */
-    node->dwork = (double *) calloc(x->size, sizeof(double));
+    node->work->dwork = (double *) calloc(x->size, sizeof(double));
 
     /* in the implementation of eval_wsum_hess we evaluate the
        child's hessian with a weight vector that has w[i] = 0
@@ -113,9 +113,7 @@ static void wsum_hess_init(expr *node)
        structural zeros, but we do not try to exploit that sparsity
        right now. */
     CSR_Matrix *Hx = x->wsum_hess;
-    node->wsum_hess = new_csr_matrix(Hx->m, Hx->n, Hx->nnz);
-    memcpy(node->wsum_hess->p, Hx->p, (Hx->m + 1) * sizeof(int));
-    memcpy(node->wsum_hess->i, Hx->i, Hx->nnz * sizeof(int));
+    node->wsum_hess = new_csr_copy_sparsity(Hx);
 }
 
 static void eval_wsum_hess(expr *node, const double *w)
@@ -126,23 +124,23 @@ static void eval_wsum_hess(expr *node, const double *w)
     if (idx->has_duplicates)
     {
         /* zero and accumulate for repeated indices */
-        memset(node->dwork, 0, x->size * sizeof(double));
+        memset(node->work->dwork, 0, x->size * sizeof(double));
         for (int i = 0; i < idx->n_idxs; i++)
         {
-            node->dwork[idx->indices[i]] += w[i];
+            node->work->dwork[idx->indices[i]] += w[i];
         }
     }
     else
     {
         /* direct write (no memset needed, no accumulation) */
         for (int i = 0; i < idx->n_idxs; i++)
         {
-            node->dwork[idx->indices[i]] = w[i];
+            node->work->dwork[idx->indices[i]] = w[i];
         }
     }
 
     /* evalute hessian of child */
-    x->eval_wsum_hess(x, node->dwork);
+    x->eval_wsum_hess(x, node->work->dwork);
     memcpy(node->wsum_hess->x, x->wsum_hess->x, x->wsum_hess->nnz * sizeof(double));
 }