From adf46d119e6e175fe890b05d8aa348cc9b6f7df3 Mon Sep 17 00:00:00 2001 From: dance858 Date: Sat, 12 Sep 2026 14:07:33 -0700 Subject: [PATCH] Migrate broadcast onto row_gather; delete the broadcast vtable slots (lesson 3, M2b) The three broadcast variants are row gathers with fixed maps over the column-major output index r = i + j*d1: ROW copies child row j, COL copies child row i, SCALAR copies row 0. The atom builds that map in jacobian_init; the broadcast_alloc / broadcast_fill_values slots, their sparse, permuted_dense and stacked_pd kernels, and the now-unused tile_int helper are removed. broadcast_type moves from matrix.h to subexpr.h, its only remaining user. Two new atom-level tests prove a pd child stays pd through ROW and COL broadcasts. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01Azb8o1AgiFRU4t9aAY7oDr --- include/subexpr.h | 8 ++ include/utils/matrix.h | 17 --- include/utils/mini_numpy.h | 1 - include/utils/permuted_dense.h | 8 -- src/atoms/affine/broadcast.c | 26 ++++- src/utils/mini_numpy.c | 8 -- src/utils/permuted_dense.c | 115 ------------------- src/utils/sparse_matrix.c | 90 --------------- src/utils/stacked_pd.c | 40 ------- tests/all_tests.c | 6 +- tests/jacobian_tests/affine/test_broadcast.h | 64 +++++++++++ tests/utils/test_permuted_dense.h | 96 ---------------- tests/utils/test_stacked_pd.h | 58 ---------- 13 files changed, 94 insertions(+), 443 deletions(-) diff --git a/include/subexpr.h b/include/subexpr.h index 58eb0c8c..b7ffd078 100644 --- a/include/subexpr.h +++ b/include/subexpr.h @@ -216,6 +216,14 @@ typedef struct index_expr bool has_duplicates; /* True if indices have duplicates (affects Hessian path) */ } index_expr; +/* Broadcast shape used by the broadcast atom. */ +typedef enum +{ + BROADCAST_ROW, /* (1, n) -> (m, n) */ + BROADCAST_COL, /* (m, 1) -> (m, n) */ + BROADCAST_SCALAR /* (1, 1) -> (m, n) */ +} broadcast_type; + typedef struct broadcast_expr { expr base; diff --git a/include/utils/matrix.h b/include/utils/matrix.h index 041f7a82..33d6bc20 100644 --- a/include/utils/matrix.h +++ b/include/utils/matrix.h @@ -23,14 +23,6 @@ #include #include -/* Broadcast shape used by the broadcast atom and its vtable methods. */ -typedef enum -{ - BROADCAST_ROW, /* (1, n) -> (m, n) */ - BROADCAST_COL, /* (m, 1) -> (m, n) */ - BROADCAST_SCALAR /* (1, 1) -> (m, n) */ -} broadcast_type; - /* Polymorphic matrix base. Concrete types embed `matrix` as their first member and implement the vtable slots below. Currently implemented: 1. sparse_matrix — generic CSR_matrix-backed matrix. @@ -89,13 +81,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); -/* 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, - int d2); -typedef void (*matrix_broadcast_fill_values_fn)(matrix *A, broadcast_type type, - int d1, int d2, matrix *out); - /* diag_vec: A is an (n, A->n) Jacobian for a length-n vector; output is (n*n, A->n) where row i lands at output row i*(n+1) (column-major diagonal positions). Other output rows are structurally zero. */ @@ -152,8 +137,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_broadcast_alloc_fn broadcast_alloc; - matrix_broadcast_fill_values_fn broadcast_fill_values; matrix_diag_vec_alloc_fn diag_vec_alloc; matrix_diag_vec_fill_values_fn diag_vec_fill_values; matrix_sum_row_partition_alloc_fn sum_row_partition_alloc; diff --git a/include/utils/mini_numpy.h b/include/utils/mini_numpy.h index d60aaabd..4998e9d8 100644 --- a/include/utils/mini_numpy.h +++ b/include/utils/mini_numpy.h @@ -25,7 +25,6 @@ void repeat(double *result, const double *a, int len, int repeats); /* Example: a = [1, 2], len = 2, tiles = 3, result = [1, 2, 1, 2, 1, 2] */ void tile_double(double *result, const double *a, int len, int tiles); -void tile_int(int *result, const int *a, int len, int tiles); /* Example: size = 5, value = 3.0, result = [3.0, 3.0, 3.0, 3.0, 3.0] */ void scaled_ones(double *result, int size, double value); diff --git a/include/utils/permuted_dense.h b/include/utils/permuted_dense.h index f8cc704d..a8f73215 100644 --- a/include/utils/permuted_dense.h +++ b/include/utils/permuted_dense.h @@ -94,14 +94,6 @@ matrix *new_permuted_dense_full(int m, int n, const double *data); place; contents are NOT preserved. */ void permuted_dense_ensure_kernel_dwork(const permuted_dense *A, size_t size); -/* Allocate C = broadcast(A, type, d1, d2), where A and C are permuted dense. */ -matrix *broadcast_pd_alloc(const permuted_dense *A, broadcast_type type, int d1, - int d2); - -/* Fill values of C = broadcast(A, type, d1, d2). */ -void broadcast_pd_fill_values(const permuted_dense *A, broadcast_type type, int d1, - int d2, permuted_dense *C); - /* Allocate C = A[map, :], where A and C are permuted dense. C stores map internally, so the fill takes none. */ matrix *row_gather_pd_alloc(const permuted_dense *A, const int *map, int m_out); diff --git a/src/atoms/affine/broadcast.c b/src/atoms/affine/broadcast.c index 1a1cccc2..fbdebc92 100644 --- a/src/atoms/affine/broadcast.c +++ b/src/atoms/affine/broadcast.c @@ -71,10 +71,25 @@ static void jacobian_init_impl(expr *node) expr *x = node->left; jacobian_init(x); - /* allocate sparsity for the broadcast output; output type matches child's. */ + /* Every output row (i, j) at column-major index i + j*d1 is a copy of one + child row: j for ROW ((1, d2) child), i for COL ((d1, 1) child), 0 for + SCALAR. A row gather with that map; the map is bound to node->jacobian + and not kept. */ broadcast_expr *bcast = (broadcast_expr *) node; - node->jacobian = - x->jacobian->broadcast_alloc(x->jacobian, bcast->type, node->d1, node->d2); + int d1 = node->d1; + int *map = (int *) sp_malloc(node->size * sizeof(int)); + for (int j = 0; j < node->d2; j++) + { + for (int i = 0; i < d1; i++) + { + int src = 0; + if (bcast->type == BROADCAST_ROW) src = j; + if (bcast->type == BROADCAST_COL) src = i; + map[i + j * d1] = src; + } + } + node->jacobian = x->jacobian->row_gather_alloc(x->jacobian, map, node->size); + sp_free(map); } static void eval_jacobian_impl(expr *node) @@ -82,9 +97,8 @@ static void eval_jacobian_impl(expr *node) eval_jacobian(node->left); /* fill values into the preallocated output. */ - broadcast_expr *bcast = (broadcast_expr *) node; - node->left->jacobian->broadcast_fill_values(node->left->jacobian, bcast->type, - node->d1, node->d2, node->jacobian); + node->left->jacobian->row_gather_fill_values(node->left->jacobian, + node->jacobian); } static void wsum_hess_init_impl(expr *node) diff --git a/src/utils/mini_numpy.c b/src/utils/mini_numpy.c index 4f436609..43c8e96b 100644 --- a/src/utils/mini_numpy.c +++ b/src/utils/mini_numpy.c @@ -38,14 +38,6 @@ void tile_double(double *result, const double *a, int len, int tiles) } } -void tile_int(int *result, const int *a, int len, int tiles) -{ - for (int i = 0; i < tiles; i++) - { - memcpy(result + i * len, a, len * sizeof(int)); - } -} - void scaled_ones(double *result, int size, double value) { for (int i = 0; i < size; i++) diff --git a/src/utils/permuted_dense.c b/src/utils/permuted_dense.c index 17958e03..637263e3 100644 --- a/src/utils/permuted_dense.c +++ b/src/utils/permuted_dense.c @@ -164,119 +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 *broadcast_pd_alloc(const permuted_dense *A, broadcast_type type, int d1, - int d2) -{ - int out_m = d1 * d2; - - int new_m0; - if (type == BROADCAST_SCALAR) - { - new_m0 = (A->m0 == 0) ? 0 : out_m; - } - else if (type == BROADCAST_ROW) - { - new_m0 = d1 * A->m0; - } - else /* BROADCAST_COL */ - { - new_m0 = d2 * A->m0; - } - - if (new_m0 == 0) - { - return new_permuted_dense(out_m, A->base.n, 0, A->n0, NULL, A->col_perm, - NULL); - } - - int *new_row_perm = (int *) sp_malloc(new_m0 * sizeof(int)); - int k = 0; - if (type == BROADCAST_SCALAR) - { - for (int i = 0; i < out_m; i++) - { - new_row_perm[k++] = i; - } - } - else if (type == BROADCAST_ROW) - { - for (int j_ii = 0; j_ii < A->m0; j_ii++) - { - int j_old = A->row_perm[j_ii]; - for (int i = 0; i < d1; i++) - { - new_row_perm[k++] = j_old * d1 + i; - } - } - } - else /* BROADCAST_COL */ - { - for (int j = 0; j < d2; j++) - { - for (int ii_old = 0; ii_old < A->m0; ii_old++) - { - new_row_perm[k++] = j * d1 + A->row_perm[ii_old]; - } - } - } - - matrix *out = new_permuted_dense(out_m, A->base.n, new_m0, A->n0, new_row_perm, - A->col_perm, NULL); - sp_free(new_row_perm); - return out; -} - -void broadcast_pd_fill_values(const permuted_dense *A, broadcast_type type, int d1, - int d2, permuted_dense *C) -{ - if (A->m0 == 0) - { - return; - } - int n0 = A->n0; - - if (type == BROADCAST_SCALAR) - { - for (int k = 0; k < C->m0; k++) - { - memcpy(C->X + k * n0, A->X, n0 * sizeof(double)); - } - } - else if (type == BROADCAST_ROW) - { - /* output row k corresponds to child dense row (k / d1). */ - (void) d2; - for (int k = 0; k < C->m0; k++) - { - memcpy(C->X + k * n0, A->X + (k / d1) * n0, n0 * sizeof(double)); - } - } - else /* BROADCAST_COL */ - { - (void) d1; - size_t child_block = A->m0 * n0; - for (int j = 0; j < d2; j++) - { - memcpy(C->X + j * child_block, A->X, child_block * sizeof(double)); - } - } -} - -static matrix *permuted_dense_vtable_broadcast_alloc(matrix *self, - broadcast_type type, int d1, - int d2) -{ - return broadcast_pd_alloc((const permuted_dense *) self, type, d1, d2); -} - -static void permuted_dense_vtable_broadcast_fill_values(matrix *self, - broadcast_type type, int d1, - int d2, matrix *out) -{ - broadcast_pd_fill_values((const permuted_dense *) self, type, d1, d2, - (permuted_dense *) out); -} - matrix *diag_vec_pd_alloc(const permuted_dense *A) { int n = A->base.m; @@ -525,8 +412,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.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; pd->base.diag_vec_fill_values = permuted_dense_vtable_diag_vec_fill_values; pd->base.sum_row_partition_alloc = permuted_dense_vtable_sum_row_partition_alloc; diff --git a/src/utils/sparse_matrix.c b/src/utils/sparse_matrix.c index b85ede9d..c1421beb 100644 --- a/src/utils/sparse_matrix.c +++ b/src/utils/sparse_matrix.c @@ -185,94 +185,6 @@ static void sparse_row_gather_fill_values(const matrix *self, matrix *out) } } -static matrix *sparse_broadcast_alloc(matrix *self, broadcast_type type, int d1, - int d2) -{ - CSR_matrix *Jx = ((sparse_matrix *) self)->csr; - int out_m = d1 * d2; - int total_nnz; - if (type == BROADCAST_ROW) - { - total_nnz = Jx->nnz * d1; - } - else if (type == BROADCAST_COL) - { - total_nnz = Jx->nnz * d2; - } - else /* BROADCAST_SCALAR */ - { - total_nnz = Jx->nnz * out_m; - } - - CSR_matrix *J = new_CSR_matrix(out_m, self->n, total_nnz); - - if (type == BROADCAST_ROW) - { - int acc = 0; - for (int i = 0; i < d2; i++) - { - int nnz_in_row = Jx->p[i + 1] - Jx->p[i]; - tile_int(J->i + acc, Jx->i + Jx->p[i], nnz_in_row, d1); - for (int rep = 0; rep < d1; rep++) - { - J->p[i * d1 + rep] = acc; - acc += nnz_in_row; - } - } - J->p[out_m] = total_nnz; - } - else if (type == BROADCAST_COL) - { - tile_int(J->i, Jx->i, Jx->nnz, d2); - int offset = 0; - for (int i = 0; i < d2; i++) - { - for (int j = 0; j < d1; j++) - { - int nnz_in_row = Jx->p[j + 1] - Jx->p[j]; - J->p[i * d1 + j] = offset; - offset += nnz_in_row; - } - } - J->p[out_m] = total_nnz; - } - else /* BROADCAST_SCALAR */ - { - tile_int(J->i, Jx->i, Jx->nnz, out_m); - int row_nnz = Jx->nnz; - for (int i = 0; i < out_m; i++) - { - J->p[i] = i * row_nnz; - } - J->p[out_m] = total_nnz; - } - return new_sparse_matrix(J); -} - -static void sparse_broadcast_fill_values(matrix *self, broadcast_type type, int d1, - int d2, matrix *out) -{ - CSR_matrix *Jx = ((sparse_matrix *) self)->csr; - if (type == BROADCAST_ROW) - { - int acc = 0; - for (int i = 0; i < d2; i++) - { - int nnz_in_row = Jx->p[i + 1] - Jx->p[i]; - tile_double(out->x + acc, Jx->x + Jx->p[i], nnz_in_row, d1); - acc += nnz_in_row * d1; - } - } - else if (type == BROADCAST_COL) - { - tile_double(out->x, Jx->x, Jx->nnz, d2); - } - else /* BROADCAST_SCALAR */ - { - tile_double(out->x, Jx->x, Jx->nnz, d1 * d2); - } -} - static matrix *sparse_diag_vec_alloc(matrix *self) { CSR_matrix *Jx = ((sparse_matrix *) self)->csr; @@ -376,8 +288,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.broadcast_alloc = sparse_broadcast_alloc; - sm->base.broadcast_fill_values = sparse_broadcast_fill_values; sm->base.diag_vec_alloc = sparse_diag_vec_alloc; sm->base.diag_vec_fill_values = sparse_diag_vec_fill_values; sm->base.sum_row_partition_alloc = sparse_sum_row_partition_alloc; diff --git a/src/utils/stacked_pd.c b/src/utils/stacked_pd.c index fca0cbed..4771d814 100644 --- a/src/utils/stacked_pd.c +++ b/src/utils/stacked_pd.c @@ -296,44 +296,6 @@ static void stacked_pd_vtable_diag_vec_fill_values(matrix *self, matrix *out) } } -// ----------------------------------------------------------------------------- -// broadcast: C = broadcast(A) where A is stacked_pd -// ----------------------------------------------------------------------------- -typedef struct -{ - broadcast_type type; - int d1; - int d2; -} pd_broadcast_ctx; - -static matrix *wrapper_pd_broadcast(permuted_dense *Bk, const void *ctx) -{ - const pd_broadcast_ctx *c = (const pd_broadcast_ctx *) ctx; - return broadcast_pd_alloc(Bk, c->type, c->d1, c->d2); -} - -static matrix *stacked_pd_vtable_broadcast_alloc(matrix *self, broadcast_type type, - int d1, int d2) -{ - stacked_pd *src = (stacked_pd *) self; - pd_broadcast_ctx ctx = {type, d1, d2}; - return spd_map_filter_blocks(src, d1 * d2, src->base.n, wrapper_pd_broadcast, - &ctx); -} - -static void stacked_pd_vtable_broadcast_fill_values(matrix *self, - broadcast_type type, int d1, - int d2, 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]; - broadcast_pd_fill_values(src->blocks[sk], type, d1, d2, out_spd->blocks[k]); - } -} - // -------------------------------------------------------------------------------- // Constructor below // -------------------------------------------------------------------------------- @@ -473,8 +435,6 @@ static void wire_vtable(stacked_pd *spd) spd->base.row_gather_fill_values = stacked_pd_vtable_row_gather_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; - spd->base.broadcast_fill_values = stacked_pd_vtable_broadcast_fill_values; spd->base.sum_row_partition_alloc = stacked_pd_vtable_sum_row_partition_alloc; } diff --git a/tests/all_tests.c b/tests/all_tests.c index fda22874..ab6235e9 100644 --- a/tests/all_tests.c +++ b/tests/all_tests.c @@ -250,6 +250,8 @@ int main(void) mu_run_test(test_broadcast_col_jacobian, tests_run); mu_run_test(test_broadcast_scalar_to_matrix_jacobian, tests_run); mu_run_test(test_double_broadcast, tests_run); + mu_run_test(test_broadcast_row_jacobian_pd_preserved, tests_run); + mu_run_test(test_broadcast_col_jacobian_pd_preserved, tests_run); mu_run_test(test_wsum_hess_multiply_1, tests_run); mu_run_test(test_wsum_hess_multiply_2, tests_run); mu_run_test(test_jacobian_trace_variable, tests_run); @@ -445,9 +447,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_broadcast_scalar, tests_run); - mu_run_test(test_permuted_dense_broadcast_row, tests_run); - mu_run_test(test_permuted_dense_broadcast_col, tests_run); mu_run_test(test_permuted_dense_diag_vec, tests_run); mu_run_test(test_permuted_dense_BTA_matching_row_perm, tests_run); mu_run_test(test_permuted_dense_BTA_empty_overlap, tests_run); @@ -560,7 +559,6 @@ int main(void) 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_diag_vec, tests_run); - mu_run_test(test_spd_vtable_broadcast_row, tests_run); mu_run_test(test_YT_kron_I, tests_run); mu_run_test(test_YT_kron_I_larger, tests_run); mu_run_test(test_I_kron_X, tests_run); diff --git a/tests/jacobian_tests/affine/test_broadcast.h b/tests/jacobian_tests/affine/test_broadcast.h index 32cfd242..c7bc101a 100644 --- a/tests/jacobian_tests/affine/test_broadcast.h +++ b/tests/jacobian_tests/affine/test_broadcast.h @@ -6,6 +6,7 @@ #include "expr.h" #include "minunit.h" #include "test_helpers.h" +#include "utils/permuted_dense.h" const char *test_broadcast_row_jacobian(void) { @@ -157,3 +158,66 @@ const char *test_double_broadcast(void) free_expr(sum); return 0; } + +/* ROW broadcast of a pd child stays pd. AU = A @ u (A 3x2 dense, u 2x1) is a + (3, 1) pd Jacobian; reshape to (1, 3) keeps it pd; broadcast to (2, 3) puts + child row j at output rows i + 2j, i.e. output row r copies A row r / 2. */ +const char *test_broadcast_row_jacobian_pd_preserved(void) +{ + double A[6] = {1.0, 2.0, 3.0, 4.0, 5.0, 6.0}; + expr *u = new_variable(2, 1, 0, 2); + expr *AU = new_left_matmul_dense(NULL, u, 3, 2, A); + expr *R = new_reshape(AU, 1, 3); + expr *B = new_broadcast(R, 2, 3); + + double u_vals[2] = {0.5, -1.5}; + jacobian_init(B); + B->forward(B, u_vals); + eval_jacobian(B); + + mu_assert("broadcast row Jacobian should be PD", B->jacobian->is_permuted_dense); + permuted_dense *pd = (permuted_dense *) B->jacobian; + mu_assert("shape", B->jacobian->m == 6 && B->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)); + for (int r = 0; r < 6; r++) + { + mu_assert("X row must be A row r / 2", + cmp_double_array(pd->X + 2 * r, A + 2 * (r / 2), 2)); + } + + free_expr(B); + return 0; +} + +/* COL broadcast of a pd child stays pd. AU = A @ u (A 2x2 dense, u 2x1) is a + (2, 1) pd Jacobian; broadcast to (2, 3) puts child row i at output rows + i + 2j, i.e. output row r copies A row r % 2. */ +const char *test_broadcast_col_jacobian_pd_preserved(void) +{ + double A[4] = {1.0, 2.0, 3.0, 4.0}; + expr *u = new_variable(2, 1, 0, 2); + expr *AU = new_left_matmul_dense(NULL, u, 2, 2, A); + expr *B = new_broadcast(AU, 2, 3); + + double u_vals[2] = {0.5, -1.5}; + jacobian_init(B); + B->forward(B, u_vals); + eval_jacobian(B); + + mu_assert("broadcast col Jacobian should be PD", B->jacobian->is_permuted_dense); + permuted_dense *pd = (permuted_dense *) B->jacobian; + mu_assert("shape", B->jacobian->m == 6 && B->jacobian->n == 2); + mu_assert("m0", pd->m0 == 6); + mu_assert("n0", pd->n0 == 2); + for (int r = 0; r < 6; r++) + { + mu_assert("X row must be A row r % 2", + cmp_double_array(pd->X + 2 * r, A + 2 * (r % 2), 2)); + } + + free_expr(B); + return 0; +} diff --git a/tests/utils/test_permuted_dense.h b/tests/utils/test_permuted_dense.h index db6d55bf..2dc6954f 100644 --- a/tests/utils/test_permuted_dense.h +++ b/tests/utils/test_permuted_dense.h @@ -387,102 +387,6 @@ const char *test_permuted_dense_row_gather(void) 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) -{ - 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); - - int d1 = 2, d2 = 3; /* out shape (2, 3), m = 6 */ - matrix *out = M->broadcast_alloc(M, BROADCAST_SCALAR, d1, d2); - permuted_dense *out_pd = (permuted_dense *) out; - - mu_assert("out m", out->m == 6); - mu_assert("out n", out->n == 5); - mu_assert("m0", out_pd->m0 == 6); - mu_assert("n0", out_pd->n0 == 2); - int expected_rp[6] = {0, 1, 2, 3, 4, 5}; - mu_assert("row_perm", cmp_int_array(out_pd->row_perm, expected_rp, 6)); - - M->broadcast_fill_values(M, BROADCAST_SCALAR, d1, d2, out); - double expected_X[12] = {7, 9, 7, 9, 7, 9, 7, 9, 7, 9, 7, 9}; - mu_assert("values", cmp_double_array(out_pd->X, expected_X, 12)); - - free_matrix(out); - free_matrix(M); - return 0; -} - -/* PD broadcast_alloc / broadcast_fill_values, ROW variant. - (1, d2) input has Jacobian of shape (d2, n_vars). Source PD: m=d2=3, - row_perm={0, 2} (rows 0 and 2 dense), col_perm={1, 4}, single dense row - per m0. Output (d1, d2) = (2, 3): each child row replicated d1=2 - times. */ -const char *test_permuted_dense_broadcast_row(void) -{ - int row_perm[2] = {0, 2}; - int col_perm[2] = {1, 4}; - double X[4] = {1.0, 2.0, /* row corresponding to child row 0 */ - 3.0, 4.0}; /* row corresponding to child row 2 */ - matrix *M = new_permuted_dense(3, 6, 2, 2, row_perm, col_perm, X); - - int d1 = 2, d2 = 3; /* output (2, 3), out m = 6 */ - matrix *out = M->broadcast_alloc(M, BROADCAST_ROW, d1, d2); - permuted_dense *out_pd = (permuted_dense *) out; - - mu_assert("out m", out->m == 6); - mu_assert("m0", out_pd->m0 == 4); /* d1 * 2 */ - mu_assert("n0", out_pd->n0 == 2); - /* row_perm = {child_row_perm[0]*d1, +1, child_row_perm[1]*d1, +1} - = {0, 1, 4, 5} */ - int expected_rp[4] = {0, 1, 4, 5}; - mu_assert("row_perm", cmp_int_array(out_pd->row_perm, expected_rp, 4)); - - M->broadcast_fill_values(M, BROADCAST_ROW, d1, d2, out); - /* each child row replicated d1 times */ - double expected_X[8] = {1.0, 2.0, 1.0, 2.0, 3.0, 4.0, 3.0, 4.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, COL variant. - (d1, 1) input has Jacobian of shape (d1, n_vars). Source PD: m=d1=3, - row_perm={0, 2}, col_perm={1, 4}, two dense rows. Output (d1, d2) = (3, 2), - out m = 6: each child row appears d2 times, shifted by j*d1. */ -const char *test_permuted_dense_broadcast_col(void) -{ - int row_perm[2] = {0, 2}; - int col_perm[2] = {1, 4}; - double X[4] = {1.0, 2.0, 3.0, 4.0}; - matrix *M = new_permuted_dense(3, 6, 2, 2, row_perm, col_perm, X); - - int d1 = 3, d2 = 2; - matrix *out = M->broadcast_alloc(M, BROADCAST_COL, d1, d2); - permuted_dense *out_pd = (permuted_dense *) out; - - mu_assert("out m", out->m == 6); - mu_assert("m0", out_pd->m0 == 4); /* d2 * 2 */ - mu_assert("n0", out_pd->n0 == 2); - /* row_perm = {0+0, 0+2, 3+0, 3+2} = {0, 2, 3, 5} */ - int expected_rp[4] = {0, 2, 3, 5}; - mu_assert("row_perm", cmp_int_array(out_pd->row_perm, expected_rp, 4)); - - M->broadcast_fill_values(M, BROADCAST_COL, d1, d2, out); - /* X = d2 copies of full source X block */ - double expected_X[8] = {1.0, 2.0, 3.0, 4.0, 1.0, 2.0, 3.0, 4.0}; - mu_assert("values", cmp_double_array(out_pd->X, expected_X, 8)); - - free_matrix(out); - free_matrix(M); - return 0; -} - /* PD diag_vec_alloc / diag_vec_fill_values. Source PD shape (3, 6) with m0=2 (rows 0 and 2) -> output PD shape (9, 6) with the same 2 dense rows mapped to positions {0, 8} = {0*4, 2*4}. */ diff --git a/tests/utils/test_stacked_pd.h b/tests/utils/test_stacked_pd.h index 47772257..cffddd56 100644 --- a/tests/utils/test_stacked_pd.h +++ b/tests/utils/test_stacked_pd.h @@ -1450,64 +1450,6 @@ const char *test_spd_vtable_diag_vec(void) return 0; } -/* broadcast_* on spd (BROADCAST_ROW): input Jac for a (1, 4) matrix has - m=4; broadcasting to (2, 4) gives output Jac with m=8. Per-block PD - rescales row_perm entries r -> {r*d1, r*d1+1, ..., r*d1+d1-1}. */ -const char *test_spd_vtable_broadcast_row(void) -{ - /* Input: 4x4 spd (Jac of a (1, 4) matrix-valued node), two blocks. - Block 0: rows {0, 1}, cols {0, 1}, X = [[1, 2], [3, 4]] - Block 1: rows {2, 3}, cols {2, 3}, X = [[5, 6], [7, 8]] */ - int row_perm_0[2] = {0, 1}; - int col_perm_0[2] = {0, 1}; - double X0[4] = {1.0, 2.0, 3.0, 4.0}; - matrix *blk0 = new_permuted_dense(4, 4, 2, 2, row_perm_0, col_perm_0, X0); - - int row_perm_1[2] = {2, 3}; - int col_perm_1[2] = {2, 3}; - double X1[4] = {5.0, 6.0, 7.0, 8.0}; - matrix *blk1 = new_permuted_dense(4, 4, 2, 2, row_perm_1, col_perm_1, X1); - - permuted_dense *blocks[2] = {(permuted_dense *) blk0, (permuted_dense *) blk1}; - matrix *M = new_stacked_pd(4, 4, 2, blocks, NULL, NULL); - - /* d1=2, d2=4 -> output Jac is 8x4 (matrix value (2, 4) vectorized). */ - matrix *C_m = M->broadcast_alloc(M, BROADCAST_ROW, 2, 4); - M->broadcast_fill_values(M, BROADCAST_ROW, 2, 4, C_m); - stacked_pd *C = (stacked_pd *) C_m; - - mu_assert("n_blocks", C->n_blocks == 2); - mu_assert("base.m", C_m->m == 8); - mu_assert("base.n", C_m->n == 4); - - /* Block 0: row_perm = {0*2, 0*2+1, 1*2, 1*2+1} = {0, 1, 2, 3}. - col_perm unchanged. X: each input row replicated d1=2 times. */ - permuted_dense *out0 = C->blocks[0]; - int expected_row_perm_0[4] = {0, 1, 2, 3}; - double expected_X0[8] = {1.0, 2.0, 1.0, 2.0, 3.0, 4.0, 3.0, 4.0}; - mu_assert("out0 m0", out0->m0 == 4); - mu_assert("out0 n0", out0->n0 == 2); - mu_assert("out0 row_perm", - cmp_int_array(out0->row_perm, expected_row_perm_0, 4)); - mu_assert("out0 col_perm", cmp_int_array(out0->col_perm, col_perm_0, 2)); - mu_assert("out0 X", cmp_double_array(out0->X, expected_X0, 8)); - - /* Block 1: row_perm = {2*2, 2*2+1, 3*2, 3*2+1} = {4, 5, 6, 7}. */ - permuted_dense *out1 = C->blocks[1]; - int expected_row_perm_1[4] = {4, 5, 6, 7}; - double expected_X1[8] = {5.0, 6.0, 5.0, 6.0, 7.0, 8.0, 7.0, 8.0}; - mu_assert("out1 m0", out1->m0 == 4); - mu_assert("out1 n0", out1->n0 == 2); - mu_assert("out1 row_perm", - cmp_int_array(out1->row_perm, expected_row_perm_1, 4)); - mu_assert("out1 col_perm", cmp_int_array(out1->col_perm, col_perm_1, 2)); - mu_assert("out1 X", cmp_double_array(out1->X, expected_X1, 8)); - - free_matrix(C_m); - free_matrix(M); - return 0; -} - /* ---------------------------------------------------------------- */ /* BA_spd_* primitive tests: C = B @ A where B is stacked_pd and A is */ /* CSC / PD / spd. Kept as siblings of the BTA_spd_* family in the */