Fix multiply's and matmul's Hessians when their children share a subexpression

src/atoms/bivariate_full_dom/matmul.c

@@ -432,8 +432,8 @@ static void wsum_hess_init_chain_rule(expr *node)
     wsum_hess_init(g);
 
     /* H = C + C^T + H_f + H_g */
-    node->wsum_hess =
-        sum_4_matrices_alloc(mnode->C, mnode->CT, f->wsum_hess, g->wsum_hess);
+    matrix *ops[4] = {mnode->C, mnode->CT, f->wsum_hess, g->wsum_hess};
+    node->wsum_hess = sum_n_matrices_alloc(ops, 4);
 
     /* allocate weight backprop workspace */
     if (!f->is_affine(f) || !g->is_affine(g))
@@ -461,23 +461,26 @@ static void eval_wsum_hess_chain_rule(expr *node, const double *w)
     BTA_matrices_fill_values(mnode->BJg, f->jacobian, mnode->C);
     transpose_matrices_fill_values(mnode->C, mnode->CT);
 
+    /* H = C + C^T + H_f + H_g, adding each term right after it is evaluated */
+    sum_n_matrices_fill_zero(node->wsum_hess);
+    sum_n_matrices_fill_add(mnode->C, 0, node->wsum_hess);
+    sum_n_matrices_fill_add(mnode->CT, 1, node->wsum_hess);
+
     /* compute Hessian of f */
     if (!is_f_affine)
     {
         Y_kron_I_vec(m, k, n, g->value, w, node->work->dwork);
         eval_wsum_hess(f, node->work->dwork);
     }
+    sum_n_matrices_fill_add(f->wsum_hess, 2, node->wsum_hess);
 
     /* compute Hessian of g */
     if (!is_g_affine)
     {
         I_kron_XT_vec(m, k, n, f->value, w, node->work->dwork);
         eval_wsum_hess(g, node->work->dwork);
     }
-
-    /* H = C + C^T + H_f + H_g */
-    sum_4_matrices_fill_values(mnode->C, mnode->CT, f->wsum_hess, g->wsum_hess,
-                               node->wsum_hess);
+    sum_n_matrices_fill_add(g->wsum_hess, 3, node->wsum_hess);
 }
 
 expr *new_matmul(expr *x, expr *y)

src/atoms/bivariate_full_dom/multiply.c

@@ -161,7 +161,8 @@ static void wsum_hess_init_impl(expr *node)
         mul_node->CT = CT;
 
         /* sparsity pattern of H = C + C^T + term2 + term3 */
-        node->wsum_hess = sum_4_matrices_alloc(C, CT, x->wsum_hess, y->wsum_hess);
+        matrix *ops[4] = {C, CT, x->wsum_hess, y->wsum_hess};
+        node->wsum_hess = sum_n_matrices_alloc(ops, 4);
     }
 }
 
@@ -189,8 +190,13 @@ static void eval_wsum_hess_impl(expr *node, const double *w)
         transpose_matrices_fill_values(mul_node->C, mul_node->CT);
 
         // ---------------------------------------------------------------
-        //              compute term2 and term 3
+        //   compute H = C + C^T + term2 + term3, adding each term right
+        //   after it is evaluated
         // ---------------------------------------------------------------
+        sum_n_matrices_fill_zero(node->wsum_hess);
+        sum_n_matrices_fill_add(mul_node->C, 0, node->wsum_hess);
+        sum_n_matrices_fill_add(mul_node->CT, 1, node->wsum_hess);
+
         if (!is_x_affine)
         {
             for (int i = 0; i < node->size; i++)
@@ -199,6 +205,7 @@ static void eval_wsum_hess_impl(expr *node, const double *w)
             }
             eval_wsum_hess(x, node->work->dwork);
         }
+        sum_n_matrices_fill_add(x->wsum_hess, 2, node->wsum_hess);
 
         if (!is_y_affine)
         {
@@ -208,12 +215,7 @@ static void eval_wsum_hess_impl(expr *node, const double *w)
             }
             eval_wsum_hess(y, node->work->dwork);
         }
-
-        // ---------------------------------------------------------------
-        //        compute H = C + C^T + term2 + term3
-        // ---------------------------------------------------------------
-        sum_4_matrices_fill_values(mul_node->C, mul_node->CT, x->wsum_hess,
-                                   y->wsum_hess, node->wsum_hess);
+        sum_n_matrices_fill_add(y->wsum_hess, 3, node->wsum_hess);
     }
 }
 

tests/all_tests.c

@@ -387,6 +387,8 @@ int main(void)
     mu_run_test(test_wsum_hess_sum_outer_product_sin_cos_left_matmul_dense,
                 tests_run);
     mu_run_test(test_wsum_hess_add_shared_subexpr, tests_run);
+    mu_run_test(test_wsum_hess_multiply_shared_subexpr, tests_run);
+    mu_run_test(test_wsum_hess_matmul_shared_subexpr, tests_run);
     mu_run_test(test_wsum_hess_matmul_exp_exp, tests_run);
     mu_run_test(test_wsum_hess_matmul_sin_cos, tests_run);
     mu_run_test(test_wsum_hess_matmul_Ax_By, tests_run);

tests/wsum_hess/composite/test_chain_rule_wsum_hess.h

@@ -636,3 +636,51 @@ const char *test_wsum_hess_add_shared_subexpr(void)
     free_expr(f);
     return 0;
 }
+
+const char *test_wsum_hess_multiply_shared_subexpr(void)
+{
+    /* multiply(log x, exp(log x)) and multiply(exp(log x), log x) with one
+     * shared log x node: evaluating the side that contains log x re-evaluates
+     * log x's Hessian with other weights, which must not leak into the other
+     * side. */
+    double u_vals[3] = {0.5, 1.0, 1.5};
+    double w[3] = {1.0, 2.0, 3.0};
+
+    expr *x = new_variable(3, 1, 0, 3);
+    expr *log_x = new_log(x);
+    expr *f = new_elementwise_mult(log_x, new_exp(log_x));
+    mu_assert("check_wsum_hess failed (shared node on the left)",
+              check_wsum_hess(f, u_vals, w, NUMERICAL_DIFF_DEFAULT_H));
+    free_expr(f);
+
+    x = new_variable(3, 1, 0, 3);
+    log_x = new_log(x);
+    f = new_elementwise_mult(new_exp(log_x), log_x);
+    mu_assert("check_wsum_hess failed (shared node on the right)",
+              check_wsum_hess(f, u_vals, w, NUMERICAL_DIFF_DEFAULT_H));
+    free_expr(f);
+    return 0;
+}
+
+const char *test_wsum_hess_matmul_shared_subexpr(void)
+{
+    /* log(X) @ exp(log(X)) and exp(log(X)) @ log(X) with one shared log(X)
+     * node, X is 2x2 */
+    double u_vals[4] = {0.5, 1.0, 1.5, 2.0};
+    double w[4] = {1.0, 2.0, 3.0, 4.0};
+
+    expr *X = new_variable(2, 2, 0, 4);
+    expr *log_X = new_log(X);
+    expr *Z = new_matmul(log_X, new_exp(log_X));
+    mu_assert("check_wsum_hess failed (shared node on the left)",
+              check_wsum_hess(Z, u_vals, w, NUMERICAL_DIFF_DEFAULT_H));
+    free_expr(Z);
+
+    X = new_variable(2, 2, 0, 4);
+    log_X = new_log(X);
+    Z = new_matmul(new_exp(log_X), log_X);
+    mu_assert("check_wsum_hess failed (shared node on the right)",
+              check_wsum_hess(Z, u_vals, w, NUMERICAL_DIFF_DEFAULT_H));
+    free_expr(Z);
+    return 0;
+}