Route convolve's Jacobian through the BA_matrices dispatchers (#157)

Author: dance858

include/subexpr.h

@@ -165,8 +165,7 @@ typedef struct convolve_expr
     expr *param_source; /* length-m kernel */
     int m;              /* kernel length */
     int n;              /* input length */
-    CSR_matrix *T;      /* (m+n-1) x n convolution matrix */
-    CSC_matrix *Jchild_CSC;
+    matrix *T;          /* (m+n-1) x n convolution matrix (sparse) */
 } convolve_expr;
 
 /* Kronecker product Z = kron(A, B) where one operand is variable-free (held by

src/atoms/affine/convolve.c

@@ -18,6 +18,7 @@
 #include "atoms/affine.h"
 #include "sparsediff_linalg/CSR_matrix.h"
 #include "sparsediff_linalg/csr_kernels.h"
+#include "sparsediff_linalg/matmul.h"
 #include "sparsediff_linalg/matrix.h"
 #include "sparsediff_linalg/tracked_alloc.h"
 #include "subexpr.h"
@@ -44,7 +45,9 @@ static void forward(expr *node, const double *u)
            which might happen when we expose Python bindings to SparseDiffEngine) */
         if (cnode->T != NULL)
         {
-            conv_matrix_fill_values(cnode->T, cnode->param_source->value);
+            conv_matrix_fill_values(cnode->T->to_csr(cnode->T),
+                                    cnode->param_source->value);
+            matrix_values_changed(cnode->T);
         }
         cnode->base.needs_parameter_refresh = false;
     }
@@ -76,15 +79,14 @@ static void jacobian_init_impl(expr *node)
     jacobian_init(child);
 
     /* Build convolution matrix of size (m+n-1) x n with m*n nonzeros */
-    cnode->T = new_CSR_matrix(m + n - 1, n, m * n);
-    conv_matrix_fill_sparsity(cnode->T, m, n);
-    conv_matrix_fill_values(cnode->T, a);
-
-    /* J = T @ J_child */
-    cnode->Jchild_CSC = csr_to_csc_alloc(child->jacobian->to_csr(child->jacobian),
-                                         node->work->iwork);
-    node->jacobian =
-        new_sparse_matrix(csr_csc_matmul_alloc(cnode->T, cnode->Jchild_CSC));
+    CSR_matrix *T = new_CSR_matrix(m + n - 1, n, m * n);
+    conv_matrix_fill_sparsity(T, m, n);
+    conv_matrix_fill_values(T, a);
+    cnode->T = new_sparse_matrix(T);
+
+    /* J = T @ J_child; its kind follows the child's (sparse, permuted_dense or
+       stacked_pd) */
+    node->jacobian = BA_matrices_alloc(cnode->T, child->jacobian);
 }
 
 static void eval_jacobian_impl(expr *node)
@@ -95,10 +97,7 @@ static void eval_jacobian_impl(expr *node)
     eval_jacobian(child);
 
     /* J = T @ J_child */
-    csr_to_csc_fill_values(child->jacobian->to_csr(child->jacobian),
-                           cnode->Jchild_CSC, node->work->iwork);
-    csr_csc_matmul_fill_values(cnode->T, cnode->Jchild_CSC,
-                               node->jacobian->to_csr(node->jacobian));
+    BA_matrices_fill_values(cnode->T, child->jacobian, node->jacobian);
 }
 
 static void wsum_hess_init_impl(expr *node)
@@ -149,8 +148,7 @@ static bool set_needs_refresh_param_source(expr *node)
 static void free_type_data(expr *node)
 {
     convolve_expr *cnode = (convolve_expr *) node;
-    free_CSR_matrix(cnode->T);
-    free_CSC_matrix(cnode->Jchild_CSC);
+    free_matrix(cnode->T);
     free_expr(cnode->param_source);
 }
 
@@ -194,9 +192,6 @@ expr *new_convolve(expr *param_node, expr *child)
     cnode->m = m;
     cnode->n = n;
 
-    /* iwork is used for csr_to_csc of child's Jacobian (size n_vars). */
-    node->work->iwork = (int *) sp_malloc(node->n_vars * sizeof(int));
-
     /* Ensure first forward() pulls current param values through any
        broadcast/promote wrappers and reflects them in T (once T is built). */
     cnode->base.needs_parameter_refresh = true;

tests/all_tests.c

@@ -266,6 +266,9 @@ int main(void)
     mu_run_test(test_jacobian_matmul, tests_run);
     mu_run_test(test_jacobian_convolve, tests_run);
     mu_run_test(test_jacobian_convolve_composite, tests_run);
+    mu_run_test(test_jacobian_convolve_csc_consumer, tests_run);
+    mu_run_test(test_jacobian_convolve_pd_child, tests_run);
+    mu_run_test(test_jacobian_convolve_spd_child, tests_run);
     mu_run_test(test_jacobian_kron_const_left, tests_run);
     mu_run_test(test_jacobian_kron_const_right, tests_run);
     mu_run_test(test_jacobian_kron_sparse, tests_run);

tests/jacobian_tests/affine/test_convolve.h

@@ -1,7 +1,9 @@
 #include <math.h>
 #include <stdio.h>
+#include <string.h>
 
 #include "atoms/affine.h"
+#include "atoms/elementwise_full_dom.h"
 #include "atoms/elementwise_restricted_dom.h"
 #include "expr.h"
 #include "minunit.h"
@@ -63,3 +65,89 @@ const char *test_jacobian_convolve_composite(void)
     free_expr(y);
     return 0;
 }
+
+const char *test_jacobian_convolve_csc_consumer(void)
+{
+    /* y = B @ conv([1, 2, 3], exp(x)) with B a sparse CSR: the left matmul
+     * (n_blocks = 1) reads conv's Jacobian through its CSC mirror. Checking at
+     * two points catches a mirror left stale by the second evaluation. */
+    double kernel[3] = {1.0, 2.0, 3.0};
+    double u1[3] = {0.1, 0.2, 0.3};
+    double u2[3] = {-0.4, 0.5, 1.0};
+
+    CSR_matrix *B = new_CSR_matrix(2, 5, 5);
+    int Bp[3] = {0, 3, 5};
+    int Bi[5] = {0, 2, 4, 1, 3};
+    double Bx[5] = {1.0, -2.0, 0.5, 3.0, 1.5};
+    memcpy(B->p, Bp, 3 * sizeof(int));
+    memcpy(B->i, Bi, 5 * sizeof(int));
+    memcpy(B->x, Bx, 5 * sizeof(double));
+
+    expr *kernel_param = new_parameter(3, 1, PARAM_FIXED, 3, kernel);
+    expr *x = new_variable(3, 1, 0, 3);
+    expr *conv = new_convolve(kernel_param, new_exp(x));
+    expr *y = new_left_matmul(NULL, conv, B);
+
+    mu_assert("check_jacobian failed at u1",
+              check_jacobian_num(y, u1, NUMERICAL_DIFF_DEFAULT_H));
+    mu_assert("check_jacobian failed at u2",
+              check_jacobian_num(y, u2, NUMERICAL_DIFF_DEFAULT_H));
+
+    free_expr(y);
+    free_CSR_matrix(B);
+    return 0;
+}
+
+const char *test_jacobian_convolve_pd_child(void)
+{
+    /* y = conv([1, 2, 3], exp(D @ x)) with D dense: the child Jacobian is
+     * permuted_dense, and so is conv's. */
+    double kernel[3] = {1.0, 2.0, 3.0};
+    double D[12] = {1.0, -0.5, 2.0, 0.5, 1.5, -1.0, 0.3, 0.2, -0.7, -1.2, 0.4, 0.9};
+    double u[3] = {0.1, -0.2, 0.3};
+
+    expr *kernel_param = new_parameter(3, 1, PARAM_FIXED, 3, kernel);
+    expr *x = new_variable(3, 1, 0, 3);
+    expr *child = new_exp(new_left_matmul_dense(NULL, x, 4, 3, D));
+    expr *y = new_convolve(kernel_param, child);
+
+    jacobian_init(y);
+    mu_assert("child jacobian must be permuted_dense",
+              child->jacobian->kind == MATRIX_PERMUTED_DENSE);
+    mu_assert("jacobian must be permuted_dense",
+              y->jacobian->kind == MATRIX_PERMUTED_DENSE);
+    mu_assert("check_jacobian failed",
+              check_jacobian_num(y, u, NUMERICAL_DIFF_DEFAULT_H));
+
+    free_expr(y);
+    return 0;
+}
+
+const char *test_jacobian_convolve_spd_child(void)
+{
+    /* y = conv([1, 2, 3], exp(vec(A @ X))) with X a 3x2 matrix variable: the
+     * dense left matmul has n_blocks = 2, so the child Jacobian is stacked_pd,
+     * and so is conv's. */
+    double kernel[3] = {1.0, 2.0, 3.0};
+    double A[6] = {1.0, -0.5, 2.0, 0.5, 1.5, -1.0};
+    double u1[6] = {0.1, 0.2, 0.3, -0.1, -0.2, -0.3};
+    double u2[6] = {0.4, -0.3, 0.2, 0.1, 0.3, -0.4};
+
+    expr *kernel_param = new_parameter(3, 1, PARAM_FIXED, 6, kernel);
+    expr *X = new_variable(3, 2, 0, 6);
+    expr *child =
+        new_exp(new_reshape(new_left_matmul_dense(NULL, X, 2, 3, A), 4, 1));
+    expr *y = new_convolve(kernel_param, child);
+
+    jacobian_init(y);
+    mu_assert("child jacobian must be stacked_pd",
+              child->jacobian->kind == MATRIX_STACKED_PD);
+    mu_assert("jacobian must be stacked_pd", y->jacobian->kind == MATRIX_STACKED_PD);
+    mu_assert("check_jacobian failed at u1",
+              check_jacobian_num(y, u1, NUMERICAL_DIFF_DEFAULT_H));
+    mu_assert("check_jacobian failed at u2",
+              check_jacobian_num(y, u2, NUMERICAL_DIFF_DEFAULT_H));
+
+    free_expr(y);
+    return 0;
+}