matmul: compute term2 through BA_kron_matrices_*

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);
 }
 

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);

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 */