[WIP] chain rule for jacobian in multiply (#64)

* function for random matrices and two tessts

* run forammter

* free matrices

* add guard for initializing jacobian or hessian twice

* run formatter

Author: dance858

include/expr.h

@@ -73,8 +73,8 @@ typedef struct expr
     CSR_Matrix *jacobian;
     CSR_Matrix *wsum_hess;
     forward_fn forward;
-    jacobian_init_fn jacobian_init;
-    wsum_hess_init_fn wsum_hess_init;
+    jacobian_init_fn jacobian_init_impl;
+    wsum_hess_init_fn wsum_hess_init_impl;
     eval_jacobian_fn eval_jacobian;
     wsum_hess_fn eval_wsum_hess;
 
@@ -99,6 +99,11 @@ void init_expr(expr *node, int d1, int d2, int n_vars, forward_fn forward,
 
 void free_expr(expr *node);
 
+/* Guarded init: skips if already initialized (safe for DAGs
+ * where a node may be visited through multiple parents). */
+void jacobian_init(expr *node);
+void wsum_hess_init(expr *node);
+
 /* Initialize CSC form of the Jacobian from the CSR Jacobian.
  * Must be called after jacobian_init. */
 void jacobian_csc_init(expr *node);

src/affine/add.c

@@ -34,11 +34,11 @@ static void forward(expr *node, const double *u)
     }
 }
 
-static void jacobian_init(expr *node)
+static void jacobian_init_impl(expr *node)
 {
     /* initialize children's jacobians */
-    node->left->jacobian_init(node->left);
-    node->right->jacobian_init(node->right);
+    jacobian_init(node->left);
+    jacobian_init(node->right);
 
     /* we never have to store more than the sum of children's nnz */
     int nnz_max = node->left->jacobian->nnz + node->right->jacobian->nnz;
@@ -60,11 +60,11 @@ static void eval_jacobian(expr *node)
                                  node->jacobian);
 }
 
-static void wsum_hess_init(expr *node)
+static void wsum_hess_init_impl(expr *node)
 {
     /* initialize children's wsum_hess */
-    node->left->wsum_hess_init(node->left);
-    node->right->wsum_hess_init(node->right);
+    wsum_hess_init(node->left);
+    wsum_hess_init(node->right);
 
     /* we never have to store more than the sum of children's nnz */
     int nnz_max = node->left->wsum_hess->nnz + node->right->wsum_hess->nnz;
@@ -95,8 +95,8 @@ expr *new_add(expr *left, expr *right)
 {
     assert(left->d1 == right->d1 && left->d2 == right->d2);
     expr *node = (expr *) calloc(1, sizeof(expr));
-    init_expr(node, left->d1, left->d2, left->n_vars, forward, jacobian_init,
-              eval_jacobian, is_affine, wsum_hess_init, eval_wsum_hess, NULL);
+    init_expr(node, left->d1, left->d2, left->n_vars, forward, jacobian_init_impl,
+              eval_jacobian, is_affine, wsum_hess_init_impl, eval_wsum_hess, NULL);
     node->left = left;
     node->right = right;
     expr_retain(left);

src/affine/broadcast.c

@@ -66,10 +66,10 @@ static void forward(expr *node, const double *u)
     }
 }
 
-static void jacobian_init(expr *node)
+static void jacobian_init_impl(expr *node)
 {
     expr *x = node->left;
-    x->jacobian_init(x);
+    jacobian_init(x);
     broadcast_expr *bcast = (broadcast_expr *) node;
     int total_nnz;
 
@@ -185,10 +185,10 @@ static void eval_jacobian(expr *node)
     }
 }
 
-static void wsum_hess_init(expr *node)
+static void wsum_hess_init_impl(expr *node)
 {
     expr *x = node->left;
-    x->wsum_hess_init(x);
+    wsum_hess_init(x);
 
     /* Same sparsity as child - weights get summed */
     node->wsum_hess = new_csr_copy_sparsity(x->wsum_hess);
@@ -279,8 +279,8 @@ expr *new_broadcast(expr *child, int d1, int d2)
     // --------------------------------------------------------------------------
     //                  initialize the rest of the expression
     // --------------------------------------------------------------------------
-    init_expr(node, d1, d2, child->n_vars, forward, jacobian_init, eval_jacobian,
-              is_affine, wsum_hess_init, eval_wsum_hess, NULL);
+    init_expr(node, d1, d2, child->n_vars, forward, jacobian_init_impl,
+              eval_jacobian, is_affine, wsum_hess_init_impl, eval_wsum_hess, NULL);
     node->left = child;
     expr_retain(child);
     bcast->type = type;

src/affine/const_scalar_mult.c

@@ -39,12 +39,12 @@ static void forward(expr *node, const double *u)
     }
 }
 
-static void jacobian_init(expr *node)
+static void jacobian_init_impl(expr *node)
 {
     expr *x = node->left;
 
     /* initialize child jacobian */
-    x->jacobian_init(x);
+    jacobian_init(x);
 
     /* same sparsity as child */
     node->jacobian = new_csr_copy_sparsity(x->jacobian);
@@ -65,12 +65,12 @@ static void eval_jacobian(expr *node)
     }
 }
 
-static void wsum_hess_init(expr *node)
+static void wsum_hess_init_impl(expr *node)
 {
     expr *x = node->left;
 
     /* initialize child's weighted Hessian */
-    x->wsum_hess_init(x);
+    wsum_hess_init(x);
 
     /* same sparsity as child */
     node->wsum_hess = new_csr_copy_sparsity(x->wsum_hess);
@@ -100,8 +100,8 @@ expr *new_const_scalar_mult(double a, expr *child)
         (const_scalar_mult_expr *) calloc(1, sizeof(const_scalar_mult_expr));
     expr *node = &mult_node->base;
 
-    init_expr(node, child->d1, child->d2, child->n_vars, forward, jacobian_init,
-              eval_jacobian, is_affine, wsum_hess_init, eval_wsum_hess, NULL);
+    init_expr(node, child->d1, child->d2, child->n_vars, forward, jacobian_init_impl,
+              eval_jacobian, is_affine, wsum_hess_init_impl, eval_wsum_hess, NULL);
     node->left = child;
     mult_node->a = a;
     expr_retain(child);

src/affine/const_vector_mult.c

@@ -38,12 +38,12 @@ static void forward(expr *node, const double *u)
     }
 }
 
-static void jacobian_init(expr *node)
+static void jacobian_init_impl(expr *node)
 {
     expr *x = node->left;
 
     /* initialize child jacobian */
-    x->jacobian_init(x);
+    jacobian_init(x);
 
     /* same sparsity as child */
     node->jacobian = new_csr_copy_sparsity(x->jacobian);
@@ -67,12 +67,12 @@ static void eval_jacobian(expr *node)
     }
 }
 
-static void wsum_hess_init(expr *node)
+static void wsum_hess_init_impl(expr *node)
 {
     expr *x = node->left;
 
     /* initialize child's weighted Hessian */
-    x->wsum_hess_init(x);
+    wsum_hess_init(x);
 
     /* same sparsity as child */
     node->wsum_hess = new_csr_copy_sparsity(x->wsum_hess);
@@ -115,8 +115,8 @@ expr *new_const_vector_mult(const double *a, expr *child)
         (const_vector_mult_expr *) calloc(1, sizeof(const_vector_mult_expr));
     expr *node = &vnode->base;
 
-    init_expr(node, child->d1, child->d2, child->n_vars, forward, jacobian_init,
-              eval_jacobian, is_affine, wsum_hess_init, eval_wsum_hess,
+    init_expr(node, child->d1, child->d2, child->n_vars, forward, jacobian_init_impl,
+              eval_jacobian, is_affine, wsum_hess_init_impl, eval_wsum_hess,
               free_type_data);
     node->left = child;
     expr_retain(child);

src/affine/constant.c

@@ -26,7 +26,7 @@ static void forward(expr *node, const double *u)
     (void) u;
 }
 
-static void jacobian_init(expr *node)
+static void jacobian_init_impl(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. */
@@ -39,7 +39,7 @@ static void eval_jacobian(expr *node)
     (void) node;
 }
 
-static void wsum_hess_init(expr *node)
+static void wsum_hess_init_impl(expr *node)
 {
     /* Constant Hessian is all zeros: n_vars x n_vars with 0 nonzeros. */
     node->wsum_hess = new_csr_matrix(node->n_vars, node->n_vars, 0);
@@ -61,8 +61,8 @@ static bool is_affine(const expr *node)
 expr *new_constant(int d1, int d2, 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);
+    init_expr(node, d1, d2, n_vars, forward, jacobian_init_impl, eval_jacobian,
+              is_affine, wsum_hess_init_impl, eval_wsum_hess, NULL);
     memcpy(node->value, values, node->size * sizeof(double));
 
     return node;

src/affine/diag_vec.c

@@ -44,11 +44,11 @@ static void forward(expr *node, const double *u)
     }
 }
 
-static void jacobian_init(expr *node)
+static void jacobian_init_impl(expr *node)
 {
     expr *x = node->left;
     int n = x->size;
-    x->jacobian_init(x);
+    jacobian_init(x);
 
     CSR_Matrix *Jx = x->jacobian;
     CSR_Matrix *J = new_csr_matrix(node->size, node->n_vars, Jx->nnz);
@@ -92,12 +92,12 @@ static void eval_jacobian(expr *node)
     }
 }
 
-static void wsum_hess_init(expr *node)
+static void wsum_hess_init_impl(expr *node)
 {
     expr *x = node->left;
 
     /* initialize child's wsum_hess */
-    x->wsum_hess_init(x);
+    wsum_hess_init(x);
 
     /* workspace for extracting diagonal weights */
     node->work->dwork = (double *) calloc(x->size, sizeof(double));
@@ -137,8 +137,8 @@ expr *new_diag_vec(expr *child)
     /* n is the number of elements (works for both row and column vectors) */
     int n = child->size;
     expr *node = (expr *) calloc(1, sizeof(expr));
-    init_expr(node, n, n, child->n_vars, forward, jacobian_init, eval_jacobian,
-              is_affine, wsum_hess_init, eval_wsum_hess, NULL);
+    init_expr(node, n, n, child->n_vars, forward, jacobian_init_impl, eval_jacobian,
+              is_affine, wsum_hess_init_impl, eval_wsum_hess, NULL);
     node->left = child;
     expr_retain(child);
 

src/affine/hstack.c

@@ -42,7 +42,7 @@ static void forward(expr *node, const double *u)
     }
 }
 
-static void jacobian_init(expr *node)
+static void jacobian_init_impl(expr *node)
 {
     hstack_expr *hnode = (hstack_expr *) node;
 
@@ -51,7 +51,7 @@ static void jacobian_init(expr *node)
     for (int i = 0; i < hnode->n_args; i++)
     {
         assert(hnode->args[i] != NULL);
-        hnode->args[i]->jacobian_init(hnode->args[i]);
+        jacobian_init(hnode->args[i]);
         nnz += hnode->args[i]->jacobian->nnz;
     }
 
@@ -100,14 +100,14 @@ static void eval_jacobian(expr *node)
     }
 }
 
-static void wsum_hess_init(expr *node)
+static void wsum_hess_init_impl(expr *node)
 {
     /* initialize children's hessians */
     hstack_expr *hnode = (hstack_expr *) node;
     int nnz = 0;
     for (int i = 0; i < hnode->n_args; i++)
     {
-        hnode->args[i]->wsum_hess_init(hnode->args[i]);
+        wsum_hess_init(hnode->args[i]);
         nnz += hnode->args[i]->wsum_hess->nnz;
     }
 
@@ -187,8 +187,9 @@ expr *new_hstack(expr **args, int n_args, int n_vars)
     /* Allocate the type-specific struct */
     hstack_expr *hnode = (hstack_expr *) calloc(1, sizeof(hstack_expr));
     expr *node = &hnode->base;
-    init_expr(node, args[0]->d1, d2, n_vars, forward, jacobian_init, eval_jacobian,
-              is_affine, wsum_hess_init, wsum_hess_eval, free_type_data);
+    init_expr(node, args[0]->d1, d2, n_vars, forward, jacobian_init_impl,
+              eval_jacobian, is_affine, wsum_hess_init_impl, wsum_hess_eval,
+              free_type_data);
 
     /* Set type-specific fields (deep copy args array) */
     hnode->args = (expr **) calloc(n_args, sizeof(expr *));

src/affine/index.c

@@ -57,11 +57,11 @@ static void forward(expr *node, const double *u)
     }
 }
 
-static void jacobian_init(expr *node)
+static void jacobian_init_impl(expr *node)
 {
     expr *x = node->left;
     index_expr *idx = (index_expr *) node;
-    x->jacobian_init(x);
+    jacobian_init(x);
 
     CSR_Matrix *Jx = x->jacobian;
     CSR_Matrix *J = new_csr_matrix(node->size, node->n_vars, Jx->nnz);
@@ -96,12 +96,12 @@ static void eval_jacobian(expr *node)
     }
 }
 
-static void wsum_hess_init(expr *node)
+static void wsum_hess_init_impl(expr *node)
 {
     expr *x = node->left;
 
     /* initialize child's wsum_hess */
-    x->wsum_hess_init(x);
+    wsum_hess_init(x);
 
     /* for setting weight vector to evaluate hessian of child */
     node->work->dwork = (double *) calloc(x->size, sizeof(double));
@@ -166,8 +166,9 @@ expr *new_index(expr *child, int d1, int d2, const int *indices, int n_idxs)
     index_expr *idx = (index_expr *) calloc(1, sizeof(index_expr));
     expr *node = &idx->base;
 
-    init_expr(node, d1, d2, child->n_vars, forward, jacobian_init, eval_jacobian,
-              is_affine, wsum_hess_init, eval_wsum_hess, free_type_data);
+    init_expr(node, d1, d2, child->n_vars, forward, jacobian_init_impl,
+              eval_jacobian, is_affine, wsum_hess_init_impl, eval_wsum_hess,
+              free_type_data);
 
     node->left = child;
     expr_retain(child);

src/affine/left_matmul.c

@@ -81,13 +81,13 @@ static void free_type_data(expr *node)
     lnode->csc_to_csr_work = NULL;
 }
 
-static void jacobian_init(expr *node)
+static void jacobian_init_impl(expr *node)
 {
     expr *x = node->left;
     left_matmul_expr *lnode = (left_matmul_expr *) node;
 
     /* initialize child's jacobian and precompute sparsity of its CSC */
-    x->jacobian_init(x);
+    jacobian_init(x);
     lnode->Jchild_CSC = csr_to_csc_fill_sparsity(x->jacobian, node->work->iwork);
 
     /* precompute sparsity of this node's jacobian in CSC and CSR */
@@ -113,11 +113,11 @@ static void eval_jacobian(expr *node)
     csc_to_csr_fill_values(J_CSC, node->jacobian, lnode->csc_to_csr_work);
 }
 
-static void wsum_hess_init(expr *node)
+static void wsum_hess_init_impl(expr *node)
 {
     /* initialize child's hessian */
     expr *x = node->left;
-    x->wsum_hess_init(x);
+    wsum_hess_init(x);
 
     /* allocate this node's hessian with the same sparsity as child's */
     node->wsum_hess = new_csr_copy_sparsity(x->wsum_hess);
@@ -169,8 +169,8 @@ expr *new_left_matmul(expr *u, const CSR_Matrix *A)
     left_matmul_expr *lnode =
         (left_matmul_expr *) calloc(1, sizeof(left_matmul_expr));
     expr *node = &lnode->base;
-    init_expr(node, d1, d2, u->n_vars, forward, jacobian_init, eval_jacobian,
-              is_affine, wsum_hess_init, eval_wsum_hess, free_type_data);
+    init_expr(node, d1, d2, u->n_vars, forward, jacobian_init_impl, eval_jacobian,
+              is_affine, wsum_hess_init_impl, eval_wsum_hess, free_type_data);
     node->left = u;
     expr_retain(u);
 
@@ -214,8 +214,8 @@ expr *new_left_matmul_dense(expr *u, int m, int n, const double *data)
     left_matmul_expr *lnode =
         (left_matmul_expr *) calloc(1, sizeof(left_matmul_expr));
     expr *node = &lnode->base;
-    init_expr(node, d1, d2, u->n_vars, forward, jacobian_init, eval_jacobian,
-              is_affine, wsum_hess_init, eval_wsum_hess, free_type_data);
+    init_expr(node, d1, d2, u->n_vars, forward, jacobian_init_impl, eval_jacobian,
+              is_affine, wsum_hess_init_impl, eval_wsum_hess, free_type_data);
     node->left = u;
     expr_retain(u);