Remove "atom specific" functions from matrix classes - PR 2

include/utils/matrix.h

@@ -89,12 +89,6 @@ typedef matrix *(*matrix_row_gather_alloc_fn)(const matrix *A, const int *map,
 /* Fill values of C = A[map, :] */
 typedef void (*matrix_row_gather_fill_values_fn)(const matrix *A, matrix *C);
 
-/* Row-tiling for the promote atom: A must be a 1-row matrix; returns
-   a new matrix of shape (size, A->n) where every row is a copy of A's
-   single row. */
-typedef matrix *(*matrix_promote_alloc_fn)(matrix *A, int size);
-typedef void (*matrix_promote_fill_values_fn)(matrix *A, matrix *out);
-
 /* Broadcast: lift the child Jacobian of a broadcast atom into the output
    Jacobian. `type` is the broadcast variant; (d1, d2) is the output shape. */
 typedef matrix *(*matrix_broadcast_alloc_fn)(matrix *A, broadcast_type type, int d1,
@@ -158,8 +152,6 @@ struct matrix
     /* Atom-specific ops */
     matrix_row_gather_alloc_fn row_gather_alloc;
     matrix_row_gather_fill_values_fn row_gather_fill_values;
-    matrix_promote_alloc_fn promote_alloc;
-    matrix_promote_fill_values_fn promote_fill_values;
     matrix_broadcast_alloc_fn broadcast_alloc;
     matrix_broadcast_fill_values_fn broadcast_fill_values;
     matrix_diag_vec_alloc_fn diag_vec_alloc;

include/utils/permuted_dense.h

@@ -109,12 +109,6 @@ matrix *row_gather_pd_alloc(const permuted_dense *A, const int *map, int m_out);
 /* Fill values of C = A[map, :], where A and C are permuted dense. */
 void row_gather_pd_fill_values(const permuted_dense *A, permuted_dense *C);
 
-/* Allocate C = promote(A, size), where A and C are permuted dense. */
-matrix *promote_pd_alloc(const permuted_dense *A, int size);
-
-/* Fill values of C = promote(A, size). */
-void promote_pd_fill_values(const permuted_dense *A, permuted_dense *C);
-
 /* Allocate C = diag_vec(A), where A and C are permuted dense. */
 matrix *diag_vec_pd_alloc(const permuted_dense *A);
 

src/atoms/affine/promote.c

@@ -40,17 +40,20 @@ static void jacobian_init_impl(expr *node)
     expr *x = node->left;
     jacobian_init(x);
 
-    /* allocate sparsity for an (node->size, n_vars) matrix whose rows are all
-       copies of the child's single row; output type matches child's type. */
-    node->jacobian = x->jacobian->promote_alloc(x->jacobian, node->size);
+    /* every output row is the child's single row (row 0): a row gather with an
+       all-zero map. The map is bound to node->jacobian, so it is not kept. */
+    int *map = (int *) sp_calloc(node->size, sizeof(int));
+    node->jacobian = x->jacobian->row_gather_alloc(x->jacobian, map, node->size);
+    sp_free(map);
 }
 
 static void eval_jacobian_impl(expr *node)
 {
     eval_jacobian(node->left);
 
     /* tile the child's single row into the preallocated output. */
-    node->left->jacobian->promote_fill_values(node->left->jacobian, node->jacobian);
+    node->left->jacobian->row_gather_fill_values(node->left->jacobian,
+                                                 node->jacobian);
 }
 
 static void wsum_hess_init_impl(expr *node)

src/utils/permuted_dense.c

@@ -164,48 +164,6 @@ static void permuted_dense_vtable_row_gather_fill_values(const matrix *self,
     row_gather_pd_fill_values((const permuted_dense *) self, (permuted_dense *) out);
 }
 
-matrix *promote_pd_alloc(const permuted_dense *A, int size)
-{
-    assert(A->m0 <= 1);
-
-    if (A->m0 == 0)
-    {
-        /* source row is all-zero; output is also structurally all-zero. */
-        return new_permuted_dense(size, A->base.n, 0, A->n0, NULL, A->col_perm,
-                                  NULL);
-    }
-
-    int *new_row_perm = (int *) sp_malloc(size * sizeof(int));
-    for (int i = 0; i < size; i++)
-    {
-        new_row_perm[i] = i;
-    }
-    matrix *out = new_permuted_dense(size, A->base.n, size, A->n0, new_row_perm,
-                                     A->col_perm, NULL);
-    sp_free(new_row_perm);
-    return out;
-}
-
-void promote_pd_fill_values(const permuted_dense *A, permuted_dense *C)
-{
-    if (A->m0 == 0) return;
-    int n0 = A->n0;
-    for (int k = 0; k < C->m0; k++)
-    {
-        memcpy(C->X + k * n0, A->X, n0 * sizeof(double));
-    }
-}
-
-static matrix *permuted_dense_vtable_promote_alloc(matrix *self, int size)
-{
-    return promote_pd_alloc((const permuted_dense *) self, size);
-}
-
-static void permuted_dense_vtable_promote_fill_values(matrix *self, matrix *out)
-{
-    promote_pd_fill_values((const permuted_dense *) self, (permuted_dense *) out);
-}
-
 matrix *broadcast_pd_alloc(const permuted_dense *A, broadcast_type type, int d1,
                            int d2)
 {
@@ -567,8 +525,6 @@ static void wire_vtable(permuted_dense *pd)
     pd->base.transpose_fill_values = permuted_dense_vtable_transpose_fill_values;
     pd->base.row_gather_alloc = permuted_dense_vtable_row_gather_alloc;
     pd->base.row_gather_fill_values = permuted_dense_vtable_row_gather_fill_values;
-    pd->base.promote_alloc = permuted_dense_vtable_promote_alloc;
-    pd->base.promote_fill_values = permuted_dense_vtable_promote_fill_values;
     pd->base.broadcast_alloc = permuted_dense_vtable_broadcast_alloc;
     pd->base.broadcast_fill_values = permuted_dense_vtable_broadcast_fill_values;
     pd->base.diag_vec_alloc = permuted_dense_vtable_diag_vec_alloc;

src/utils/sparse_matrix.c

@@ -185,32 +185,6 @@ static void sparse_row_gather_fill_values(const matrix *self, matrix *out)
     }
 }
 
-static matrix *sparse_promote_alloc(matrix *self, int size)
-{
-    CSR_matrix *Jx = ((sparse_matrix *) self)->csr;
-    int row_nnz = Jx->nnz;
-    CSR_matrix *J = new_CSR_matrix(size, self->n, size * row_nnz);
-
-    for (int row = 0; row < size; row++)
-    {
-        J->p[row] = row * row_nnz;
-        memcpy(J->i + row * row_nnz, Jx->i, row_nnz * sizeof(int));
-    }
-    J->p[size] = size * row_nnz;
-    J->nnz = size * row_nnz;
-    return new_sparse_matrix(J);
-}
-
-static void sparse_promote_fill_values(matrix *self, matrix *out)
-{
-    CSR_matrix *Jx = ((sparse_matrix *) self)->csr;
-    int row_nnz = Jx->nnz;
-    for (int row = 0; row < out->m; row++)
-    {
-        memcpy(out->x + row * row_nnz, Jx->x, row_nnz * sizeof(double));
-    }
-}
-
 static matrix *sparse_broadcast_alloc(matrix *self, broadcast_type type, int d1,
                                       int d2)
 {
@@ -402,8 +376,6 @@ static void wire_vtable(sparse_matrix *sm)
     sm->base.transpose_fill_values = sparse_transpose_fill_values;
     sm->base.row_gather_alloc = sparse_row_gather_alloc;
     sm->base.row_gather_fill_values = sparse_row_gather_fill_values;
-    sm->base.promote_alloc = sparse_promote_alloc;
-    sm->base.promote_fill_values = sparse_promote_fill_values;
     sm->base.broadcast_alloc = sparse_broadcast_alloc;
     sm->base.broadcast_fill_values = sparse_broadcast_fill_values;
     sm->base.diag_vec_alloc = sparse_diag_vec_alloc;

src/utils/stacked_pd.c

@@ -267,31 +267,6 @@ static void stacked_pd_vtable_row_gather_fill_values(const matrix *self, matrix
     }
 }
 
-// -----------------------------------------------------------------------------
-//        promote: C = promote(A) where A is stacked_pd
-// -----------------------------------------------------------------------------
-static matrix *wrapper_pd_promote(permuted_dense *Bk, const void *ctx)
-{
-    return promote_pd_alloc(Bk, *(const int *) ctx);
-}
-
-static matrix *stacked_pd_vtable_promote_alloc(matrix *self, int size)
-{
-    stacked_pd *src = (stacked_pd *) self;
-    return spd_map_filter_blocks(src, size, src->base.n, wrapper_pd_promote, &size);
-}
-
-static void stacked_pd_vtable_promote_fill_values(matrix *self, matrix *out)
-{
-    stacked_pd *src = (stacked_pd *) self;
-    stacked_pd *out_spd = (stacked_pd *) out;
-    for (int k = 0; k < out_spd->n_blocks; k++)
-    {
-        int sk = out_spd->src_block_idx[k];
-        promote_pd_fill_values(src->blocks[sk], out_spd->blocks[k]);
-    }
-}
-
 // -----------------------------------------------------------------------------
 //        diag_vec: C = diag(vec(A)) where A is stacked_pd
 // -----------------------------------------------------------------------------
@@ -496,8 +471,6 @@ static void wire_vtable(stacked_pd *spd)
     spd->base.to_csr = stacked_pd_to_csr;
     spd->base.row_gather_alloc = stacked_pd_vtable_row_gather_alloc;
     spd->base.row_gather_fill_values = stacked_pd_vtable_row_gather_fill_values;
-    spd->base.promote_alloc = stacked_pd_vtable_promote_alloc;
-    spd->base.promote_fill_values = stacked_pd_vtable_promote_fill_values;
     spd->base.diag_vec_alloc = stacked_pd_vtable_diag_vec_alloc;
     spd->base.diag_vec_fill_values = stacked_pd_vtable_diag_vec_fill_values;
     spd->base.broadcast_alloc = stacked_pd_vtable_broadcast_alloc;

tests/all_tests.c

@@ -245,6 +245,7 @@ int main(void)
     mu_run_test(test_sum_of_index, tests_run);
     mu_run_test(test_promote_scalar_jacobian, tests_run);
     mu_run_test(test_promote_scalar_to_matrix_jacobian, tests_run);
+    mu_run_test(test_promote_jacobian_pd_preserved, tests_run);
     mu_run_test(test_broadcast_row_jacobian, tests_run);
     mu_run_test(test_broadcast_col_jacobian, tests_run);
     mu_run_test(test_broadcast_scalar_to_matrix_jacobian, tests_run);
@@ -444,7 +445,6 @@ int main(void)
 #ifdef SP_TRACK_MEMORY
     mu_run_test(test_row_gather_spd_fill_no_transient_alloc, tests_run);
 #endif
-    mu_run_test(test_permuted_dense_promote, tests_run);
     mu_run_test(test_permuted_dense_broadcast_scalar, tests_run);
     mu_run_test(test_permuted_dense_broadcast_row, tests_run);
     mu_run_test(test_permuted_dense_broadcast_col, tests_run);
@@ -559,7 +559,6 @@ int main(void)
     mu_run_test(test_spd_vtable_transpose, tests_run);
     mu_run_test(test_spd_vtable_refresh_csc_values_noop, tests_run);
     mu_run_test(test_spd_vtable_row_gather, tests_run);
-    mu_run_test(test_spd_vtable_promote, tests_run);
     mu_run_test(test_spd_vtable_diag_vec, tests_run);
     mu_run_test(test_spd_vtable_broadcast_row, tests_run);
     mu_run_test(test_YT_kron_I, tests_run);

tests/jacobian_tests/affine/test_promote.h

@@ -6,6 +6,7 @@
 #include "expr.h"
 #include "minunit.h"
 #include "test_helpers.h"
+#include "utils/permuted_dense.h"
 
 const char *test_promote_scalar_jacobian(void)
 {
@@ -58,3 +59,33 @@ const char *test_promote_scalar_to_matrix_jacobian(void)
     free_expr(promote_node);
     return 0;
 }
+
+/* A pd child Jacobian stays pd through promote: AU = A @ u with A a 1x2 dense
+   constant and u a 2x1 variable has a (1, 2) pd Jacobian with m0 = 1;
+   promote(AU, 3, 2) gathers that single row six times. */
+const char *test_promote_jacobian_pd_preserved(void)
+{
+    double A[2] = {1.5, -2.0};
+    expr *u = new_variable(2, 1, 0, 2);
+    expr *AU = new_left_matmul_dense(NULL, u, 1, 2, A);
+    expr *P = new_promote(AU, 3, 2);
+
+    double u_vals[2] = {0.5, -1.5};
+    jacobian_init(P);
+    P->forward(P, u_vals);
+    eval_jacobian(P);
+
+    mu_assert("promote Jacobian should be PD", P->jacobian->is_permuted_dense);
+    permuted_dense *pd = (permuted_dense *) P->jacobian;
+    mu_assert("shape", P->jacobian->m == 6 && P->jacobian->n == 2);
+    mu_assert("m0", pd->m0 == 6);
+    mu_assert("n0", pd->n0 == 2);
+    int expected_row_perm[6] = {0, 1, 2, 3, 4, 5};
+    mu_assert("row_perm", cmp_int_array(pd->row_perm, expected_row_perm, 6));
+    double expected_X[12] = {1.5, -2.0, 1.5, -2.0, 1.5, -2.0,
+                             1.5, -2.0, 1.5, -2.0, 1.5, -2.0};
+    mu_assert("X values", cmp_double_array(pd->X, expected_X, 12));
+
+    free_expr(P);
+    return 0;
+}

tests/utils/test_permuted_dense.h

@@ -387,40 +387,6 @@ const char *test_permuted_dense_row_gather(void)
     return 0;
 }
 
-/* PD promote_alloc / promote_fill_values: tile a 1-row PD into a
-   `size`-row PD where every row is a copy of the source row. */
-const char *test_permuted_dense_promote(void)
-{
-    /* Source PD, shape (1, 5), single dense row at row 0, cols {1, 3}. */
-    int row_perm[1] = {0};
-    int col_perm[2] = {1, 3};
-    double X[2] = {7.0, 9.0};
-    matrix *M = new_permuted_dense(1, 5, 1, 2, row_perm, col_perm, X);
-
-    matrix *out = M->promote_alloc(M, 4);
-    permuted_dense *out_pd = (permuted_dense *) out;
-
-    mu_assert("out m", out->m == 4);
-    mu_assert("out n", out->n == 5);
-    mu_assert("out nnz", out->nnz == 8); /* m0=4 * n0=2 */
-    mu_assert("m0", out_pd->m0 == 4);
-    mu_assert("n0", out_pd->n0 == 2);
-
-    int expected_row_perm[4] = {0, 1, 2, 3};
-    mu_assert("row_perm", cmp_int_array(out_pd->row_perm, expected_row_perm, 4));
-    int expected_col_perm[2] = {1, 3};
-    mu_assert("col_perm", cmp_int_array(out_pd->col_perm, expected_col_perm, 2));
-
-    M->promote_fill_values(M, out);
-
-    double expected_X[8] = {7.0, 9.0, 7.0, 9.0, 7.0, 9.0, 7.0, 9.0};
-    mu_assert("values", cmp_double_array(out_pd->X, expected_X, 8));
-
-    free_matrix(out);
-    free_matrix(M);
-    return 0;
-}
-
 /* PD broadcast_alloc / broadcast_fill_values, SCALAR variant.
    (1, 5) PD with single dense row -> (d1*d2, 5) PD with that row tiled. */
 const char *test_permuted_dense_broadcast_scalar(void)

tests/utils/test_stacked_pd.h

@@ -1395,43 +1395,6 @@ const char *test_spd_vtable_row_gather(void)
     return 0;
 }
 
-/* promote_* on spd: replicate the single row across `size` rows; per-block
-   delegation drops empty blocks so output has at most one block. */
-const char *test_spd_vtable_promote(void)
-{
-    /* 1x4 spd, single block carries row 0 at cols {0, 2} with values
-       [9, 11]. Promote to size=3.                                         */
-    int row_perm[1] = {0};
-    int col_perm[2] = {0, 2};
-    double X[2] = {9.0, 11.0};
-    matrix *blk = new_permuted_dense(1, 4, 1, 2, row_perm, col_perm, X);
-
-    permuted_dense *blocks[1] = {(permuted_dense *) blk};
-    matrix *M = new_stacked_pd(1, 4, 1, blocks, NULL, NULL);
-
-    matrix *C_m = M->promote_alloc(M, 3);
-    M->promote_fill_values(M, C_m);
-    stacked_pd *C = (stacked_pd *) C_m;
-
-    mu_assert("n_blocks", C->n_blocks == 1);
-    mu_assert("base.m", C_m->m == 3);
-    mu_assert("base.n", C_m->n == 4);
-
-    permuted_dense *out0 = C->blocks[0];
-    int expected_row_perm[3] = {0, 1, 2};
-    int expected_col_perm[2] = {0, 2};
-    double expected_X[6] = {9.0, 11.0, 9.0, 11.0, 9.0, 11.0};
-    mu_assert("out0 m0", out0->m0 == 3);
-    mu_assert("out0 n0", out0->n0 == 2);
-    mu_assert("out0 row_perm", cmp_int_array(out0->row_perm, expected_row_perm, 3));
-    mu_assert("out0 col_perm", cmp_int_array(out0->col_perm, expected_col_perm, 2));
-    mu_assert("out0 X", cmp_double_array(out0->X, expected_X, 6));
-
-    free_matrix(C_m);
-    free_matrix(M);
-    return 0;
-}
-
 /* diag_vec_* on spd: per-block row_perm entries are rescaled r -> r*(n+1);
    X buffers are unchanged; structure is preserved (same n_blocks). */
 const char *test_spd_vtable_diag_vec(void)