From fdb658680f1174fb49796a9860201e8086cb7727 Mon Sep 17 00:00:00 2001 From: dance858 Date: Tue, 13 Oct 2026 17:59:37 +0200 Subject: [PATCH] matmul: compute term2 through BA_kron_matrices_* term2 = kron(I_n, X) @ J_g only feeds sum_matrices_*, which takes any kind, so it no longer needs BA_dense_kron_matrices_* directly. For n == 1 the dispatcher routes through BA_matrices_* and term2 is permuted_dense. This is the last direct engine use of BA_dense_kron_matrices_*, so the library can make it private. Adds test_jacobian_matmul_exp_X_exp_y, the first chain-rule matmul test with n == 1: checks the Jacobian against finite differences at two points and that term2 is permuted_dense. Co-Authored-By: Claude Opus 5.5 --- src/atoms/bivariate_full_dom/matmul.c | 11 +++++---- tests/all_tests.c | 1 + .../composite/test_chain_rule_jacobian.h | 23 +++++++++++++++++++ 3 files changed, 30 insertions(+), 5 deletions(-) diff --git a/src/atoms/bivariate_full_dom/matmul.c b/src/atoms/bivariate_full_dom/matmul.c index 43493f6..7f3fa79 100644 --- a/src/atoms/bivariate_full_dom/matmul.c +++ b/src/atoms/bivariate_full_dom/matmul.c @@ -234,7 +234,7 @@ static void jacobian_init_chain_rule(expr *node) mnode->YT = new_permuted_dense_full(n, k, NULL); mnode->X = new_permuted_dense_full(m, k, NULL); mnode->term1 = BA_dense_kron_I_matrices_alloc(mnode->YT, m, f->jacobian); - mnode->term2 = BA_dense_kron_matrices_alloc(mnode->X, n, g->jacobian); + mnode->term2 = BA_kron_matrices_alloc(mnode->X, n, g->jacobian); int max_nnz = mnode->term1->nnz + mnode->term2->nnz; max_nnz = MIN(max_nnz, sat_mul_int(node->size, node->n_vars)); node->jacobian = new_sparse_matrix_alloc(node->size, node->n_vars, max_nnz); @@ -261,11 +261,12 @@ static void eval_jacobian_chain_rule(expr *node) A_transpose(mnode->X->x, f->value, k, m); matrix_values_changed(mnode->X); - /* evaluate term1, term2, and their sum. The sum reads the two stacked_pd - terms through their CSR views, one O(nnz) refresh per term per eval; a - bound-map sum in the library would remove that. */ + /* evaluate term1, term2, and their sum. term1 is stacked_pd and term2 is + stacked_pd for n > 1 (permuted_dense for n == 1); the sum reads them + through their CSR views, one O(nnz) refresh per stacked_pd term per + eval; a bound-map sum in the library would remove that. */ BA_dense_kron_I_matrices_fill_values(mnode->YT, m, f->jacobian, mnode->term1); - BA_dense_kron_matrices_fill_values(mnode->X, n, g->jacobian, mnode->term2); + BA_kron_matrices_fill_values(mnode->X, n, g->jacobian, mnode->term2); sum_matrices_fill_values(mnode->term1, mnode->term2, node->jacobian); } diff --git a/tests/all_tests.c b/tests/all_tests.c index 55bd587..704383d 100644 --- a/tests/all_tests.c +++ b/tests/all_tests.c @@ -167,6 +167,7 @@ int main(void) mu_run_test(test_jacobian_atan_Ax, tests_run); mu_run_test(test_jacobian_quad_form_exp, tests_run); mu_run_test(test_jacobian_matmul_exp_exp, tests_run); + mu_run_test(test_jacobian_matmul_exp_X_exp_y, tests_run); mu_run_test(test_jacobian_matmul_sin_cos, tests_run); mu_run_test(test_jacobian_matmul_Ax_By, tests_run); mu_run_test(test_jacobian_matmul_sin_Ax_cos_Bx, tests_run); diff --git a/tests/jacobian_tests/composite/test_chain_rule_jacobian.h b/tests/jacobian_tests/composite/test_chain_rule_jacobian.h index 446aca7..8f0781c 100644 --- a/tests/jacobian_tests/composite/test_chain_rule_jacobian.h +++ b/tests/jacobian_tests/composite/test_chain_rule_jacobian.h @@ -209,6 +209,29 @@ const char *test_jacobian_matmul_exp_exp(void) return 0; } +/* Z = exp(X) @ exp(y) with y a column (n == 1): term2 = kron(I_1, X) @ J_g goes + through BA_matrices_* and is permuted_dense. */ +const char *test_jacobian_matmul_exp_X_exp_y(void) +{ + /* X is 2x3 (vars 0-5), y is 3x1 (vars 6-8) */ + double u1[9] = {0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9}; + double u2[9] = {-0.4, 0.3, 1.1, -0.2, 0.6, 0.05, -0.7, 0.9, 0.25}; + + expr *X = new_variable(2, 3, 0, 9); + expr *y = new_variable(3, 1, 6, 9); + expr *Z = new_matmul(new_exp(X), new_exp(y)); + + mu_assert("check_jacobian failed at u1", + check_jacobian_num(Z, u1, NUMERICAL_DIFF_DEFAULT_H)); + mu_assert("check_jacobian failed at u2", + check_jacobian_num(Z, u2, NUMERICAL_DIFF_DEFAULT_H)); + mu_assert("term2 should be permuted_dense for n == 1", + ((matmul_expr *) Z)->term2->kind == MATRIX_PERMUTED_DENSE); + + free_expr(Z); + return 0; +} + const char *test_jacobian_matmul_sin_cos(void) { /* Z = sin(X) @ cos(Y), X is 2x2, Y is 2x3 */