diff --git a/src/atoms/bivariate_full_dom/matmul.c b/src/atoms/bivariate_full_dom/matmul.c index 9e4e681..b70f304 100644 --- a/src/atoms/bivariate_full_dom/matmul.c +++ b/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,12 +461,18 @@ 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) @@ -474,10 +480,7 @@ static void eval_wsum_hess_chain_rule(expr *node, const double *w) 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) diff --git a/src/atoms/bivariate_full_dom/multiply.c b/src/atoms/bivariate_full_dom/multiply.c index d2c37fc..8067a15 100644 --- a/src/atoms/bivariate_full_dom/multiply.c +++ b/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); } } diff --git a/tests/all_tests.c b/tests/all_tests.c index 704383d..6131b8b 100644 --- a/tests/all_tests.c +++ b/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); diff --git a/tests/wsum_hess/composite/test_chain_rule_wsum_hess.h b/tests/wsum_hess/composite/test_chain_rule_wsum_hess.h index 5f9ecb4..06b4d0c 100644 --- a/tests/wsum_hess/composite/test_chain_rule_wsum_hess.h +++ b/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; +}