From f30f6f174c4c5bc13ef01220a651d59b3036387c Mon Sep 17 00:00:00 2001 From: William Zijie Zhang Date: Wed, 18 Mar 2026 12:35:50 -0400 Subject: [PATCH 1/3] Add vstack atom implementation plan Co-Authored-By: Claude Opus 4.6 (1M context) --- docs/vstack-plan.md | 112 ++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 112 insertions(+) create mode 100644 docs/vstack-plan.md diff --git a/docs/vstack-plan.md b/docs/vstack-plan.md new file mode 100644 index 00000000..a85871a0 --- /dev/null +++ b/docs/vstack-plan.md @@ -0,0 +1,112 @@ +# Plan: Implement `vstack` atom + +## Context + +The engine needs a `vstack` (vertical stack) operation that concatenates expressions along the row dimension. Unlike `hstack` — where column-major storage means output values are a simple flat concatenation of children — `vstack` interleaves children's data across columns, making every operation (forward, Jacobian, Hessian) require row-level remapping. + +## Core challenge: column-major interleaving + +For `vstack` of children with shapes `(r_0, d2), (r_1, d2), ...`: +- Output shape: `(R, d2)` where `R = sum(r_i)` +- Column `j` of the output = `[child_0 col j, child_1 col j, ...]` stacked +- Output flat index `k` maps to: `global_row = k % R`, `col = k / R` +- Finding which child owns `global_row` requires a precomputed lookup + +## Implementation + +### 1. Type-specific struct (`include/subexpr.h`) + +```c +typedef struct vstack_expr +{ + expr base; + expr **args; + int n_args; + int *row_offsets; /* row_offsets[i] = sum of args[0..i-1]->d1 */ + int *row_to_child; /* length R: maps global row -> child index */ + int *row_to_local; /* length R: maps global row -> local row in child */ + CSR_Matrix *CSR_work; /* for Hessian accumulation */ +} vstack_expr; +``` + +Precompute `row_to_child[R]` and `row_to_local[R]` once in the constructor. These let every later operation do O(1) lookups instead of scanning children. + +### 2. Constructor (`src/affine/vstack.c` — `new_vstack`) + +- Assert all children have same `d2` +- Compute `R = sum(r_i)`, output is `(R, d2)` +- Allocate `vstack_expr`, call `init_expr` +- Build `row_offsets`, `row_to_child`, `row_to_local` arrays +- Retain all children + +### 3. Forward pass + +Loop over `d2` columns; within each column, loop over children and `memcpy` each child's column slice: + +``` +offset = 0 +for j in 0..d2-1: + for i in 0..n_args-1: + memcpy(value + offset, args[i]->value + j * r_i, r_i * sizeof(double)) + offset += r_i +``` + +### 4. Jacobian init + +The output Jacobian is `(size, n_vars)` in CSR. Each output row `k` corresponds to a specific child Jacobian row. Iterate output rows in order, use the precomputed mapping to copy column indices from the right child row: + +``` +for k in 0..size-1: + global_row = k % R + col = k / R + child_idx = row_to_child[global_row] + local_row = row_to_local[global_row] + child_flat = col * args[child_idx]->d1 + local_row + copy column indices from child's Jacobian row child_flat + set J->p[k+1] +``` + +### 5. eval_jacobian + +Same loop as jacobian_init but copies values (`x`) instead of column indices (`i`). + +### 6. Hessian (wsum_hess_init / eval) + +**Sparsity init**: Identical to hstack — iteratively union children's Hessian patterns using `sum_csr_matrices_fill_sparsity`. + +**Eval**: The weight vector `w` has one entry per output element. Each child needs a gathered weight vector. Allocate `dwork` of size `max(child->size)`. For each child, gather its weights: + +``` +for j in 0..d2-1: + for r in 0..r_i-1: + dwork[j * r_i + r] = w[j * R + row_offsets[i] + r] +``` + +Then call `child->eval_wsum_hess(child, dwork)` and accumulate with `sum_csr_matrices_fill_values`. This follows the same pattern transpose uses for weight permutation (`transpose.c:111-117`). + +### 7. is_affine / free_type_data + +Same pattern as hstack — check all children / free all children + auxiliary arrays. + +## Files to modify + +| File | Change | +|------|--------| +| `include/subexpr.h` | Add `vstack_expr` struct | +| `include/affine.h` | Declare `new_vstack` | +| `src/affine/vstack.c` | New file — full implementation | +| `src/affine/hstack.c` | Reference only (pattern to follow) | +| `src/affine/transpose.c` | Reference only (weight permutation pattern) | +| `tests/forward_pass/affine/test_vstack.h` | New — forward pass tests | +| `tests/jacobian_tests/test_vstack.h` | New — Jacobian tests | +| `tests/wsum_hess/test_vstack.h` | New — Hessian tests | +| `tests/all_tests.c` | Include new test headers, add `mu_run_test` calls | + +## Verification + +1. `cmake -B build -S . -DCMAKE_BUILD_TYPE=Debug && cmake --build build` +2. `./build/all_tests` — all existing + new vstack tests pass +3. Test cases should cover: + - Vectors: `vstack([log(x), exp(x)])` where x is `(3,1)` — output `(6,1)` (degenerate case: behaves like hstack for vectors) + - Matrices: `vstack([log(x), exp(y)])` where x is `(2,3)`, y is `(1,3)` — output `(3,3)` — verifies interleaving + - Hessian with shared variables: `vstack([log(x), exp(x)])` — verifies weight gathering and Hessian accumulation From 0a6a4b1a7d8d646d55956750661e062ead059286 Mon Sep 17 00:00:00 2001 From: William Zijie Zhang Date: Wed, 25 Mar 2026 10:49:01 -0400 Subject: [PATCH 2/3] Implement vstack atom via transpose(hstack(transpose(...))) composition Reuses existing hstack and transpose implementations instead of a from-scratch approach, avoiding manual row-interleaving logic. Removes the original vstack-plan.md since the composition strategy supersedes it. Co-Authored-By: Claude Opus 4.6 (1M context) --- docs/vstack-plan.md | 112 ----------------------- include/affine.h | 1 + src/affine/vstack.c | 47 ++++++++++ tests/all_tests.c | 9 ++ tests/forward_pass/affine/test_vstack.h | 87 ++++++++++++++++++ tests/jacobian_tests/test_vstack.h | 99 +++++++++++++++++++++ tests/wsum_hess/test_vstack.h | 113 ++++++++++++++++++++++++ 7 files changed, 356 insertions(+), 112 deletions(-) delete mode 100644 docs/vstack-plan.md create mode 100644 src/affine/vstack.c create mode 100644 tests/forward_pass/affine/test_vstack.h create mode 100644 tests/jacobian_tests/test_vstack.h create mode 100644 tests/wsum_hess/test_vstack.h diff --git a/docs/vstack-plan.md b/docs/vstack-plan.md deleted file mode 100644 index a85871a0..00000000 --- a/docs/vstack-plan.md +++ /dev/null @@ -1,112 +0,0 @@ -# Plan: Implement `vstack` atom - -## Context - -The engine needs a `vstack` (vertical stack) operation that concatenates expressions along the row dimension. Unlike `hstack` — where column-major storage means output values are a simple flat concatenation of children — `vstack` interleaves children's data across columns, making every operation (forward, Jacobian, Hessian) require row-level remapping. - -## Core challenge: column-major interleaving - -For `vstack` of children with shapes `(r_0, d2), (r_1, d2), ...`: -- Output shape: `(R, d2)` where `R = sum(r_i)` -- Column `j` of the output = `[child_0 col j, child_1 col j, ...]` stacked -- Output flat index `k` maps to: `global_row = k % R`, `col = k / R` -- Finding which child owns `global_row` requires a precomputed lookup - -## Implementation - -### 1. Type-specific struct (`include/subexpr.h`) - -```c -typedef struct vstack_expr -{ - expr base; - expr **args; - int n_args; - int *row_offsets; /* row_offsets[i] = sum of args[0..i-1]->d1 */ - int *row_to_child; /* length R: maps global row -> child index */ - int *row_to_local; /* length R: maps global row -> local row in child */ - CSR_Matrix *CSR_work; /* for Hessian accumulation */ -} vstack_expr; -``` - -Precompute `row_to_child[R]` and `row_to_local[R]` once in the constructor. These let every later operation do O(1) lookups instead of scanning children. - -### 2. Constructor (`src/affine/vstack.c` — `new_vstack`) - -- Assert all children have same `d2` -- Compute `R = sum(r_i)`, output is `(R, d2)` -- Allocate `vstack_expr`, call `init_expr` -- Build `row_offsets`, `row_to_child`, `row_to_local` arrays -- Retain all children - -### 3. Forward pass - -Loop over `d2` columns; within each column, loop over children and `memcpy` each child's column slice: - -``` -offset = 0 -for j in 0..d2-1: - for i in 0..n_args-1: - memcpy(value + offset, args[i]->value + j * r_i, r_i * sizeof(double)) - offset += r_i -``` - -### 4. Jacobian init - -The output Jacobian is `(size, n_vars)` in CSR. Each output row `k` corresponds to a specific child Jacobian row. Iterate output rows in order, use the precomputed mapping to copy column indices from the right child row: - -``` -for k in 0..size-1: - global_row = k % R - col = k / R - child_idx = row_to_child[global_row] - local_row = row_to_local[global_row] - child_flat = col * args[child_idx]->d1 + local_row - copy column indices from child's Jacobian row child_flat - set J->p[k+1] -``` - -### 5. eval_jacobian - -Same loop as jacobian_init but copies values (`x`) instead of column indices (`i`). - -### 6. Hessian (wsum_hess_init / eval) - -**Sparsity init**: Identical to hstack — iteratively union children's Hessian patterns using `sum_csr_matrices_fill_sparsity`. - -**Eval**: The weight vector `w` has one entry per output element. Each child needs a gathered weight vector. Allocate `dwork` of size `max(child->size)`. For each child, gather its weights: - -``` -for j in 0..d2-1: - for r in 0..r_i-1: - dwork[j * r_i + r] = w[j * R + row_offsets[i] + r] -``` - -Then call `child->eval_wsum_hess(child, dwork)` and accumulate with `sum_csr_matrices_fill_values`. This follows the same pattern transpose uses for weight permutation (`transpose.c:111-117`). - -### 7. is_affine / free_type_data - -Same pattern as hstack — check all children / free all children + auxiliary arrays. - -## Files to modify - -| File | Change | -|------|--------| -| `include/subexpr.h` | Add `vstack_expr` struct | -| `include/affine.h` | Declare `new_vstack` | -| `src/affine/vstack.c` | New file — full implementation | -| `src/affine/hstack.c` | Reference only (pattern to follow) | -| `src/affine/transpose.c` | Reference only (weight permutation pattern) | -| `tests/forward_pass/affine/test_vstack.h` | New — forward pass tests | -| `tests/jacobian_tests/test_vstack.h` | New — Jacobian tests | -| `tests/wsum_hess/test_vstack.h` | New — Hessian tests | -| `tests/all_tests.c` | Include new test headers, add `mu_run_test` calls | - -## Verification - -1. `cmake -B build -S . -DCMAKE_BUILD_TYPE=Debug && cmake --build build` -2. `./build/all_tests` — all existing + new vstack tests pass -3. Test cases should cover: - - Vectors: `vstack([log(x), exp(x)])` where x is `(3,1)` — output `(6,1)` (degenerate case: behaves like hstack for vectors) - - Matrices: `vstack([log(x), exp(y)])` where x is `(2,3)`, y is `(1,3)` — output `(3,3)` — verifies interleaving - - Hessian with shared variables: `vstack([log(x), exp(x)])` — verifies weight gathering and Hessian accumulation diff --git a/include/affine.h b/include/affine.h index ada58f1f..eccd5fbb 100644 --- a/include/affine.h +++ b/include/affine.h @@ -29,6 +29,7 @@ expr *new_neg(expr *child); expr *new_sum(expr *child, int axis); expr *new_hstack(expr **args, int n_args, int n_vars); +expr *new_vstack(expr **args, int n_args, int n_vars); expr *new_promote(expr *child, int d1, int d2); expr *new_trace(expr *child); diff --git a/src/affine/vstack.c b/src/affine/vstack.c new file mode 100644 index 00000000..6fcfbd0b --- /dev/null +++ b/src/affine/vstack.c @@ -0,0 +1,47 @@ +/* + * Copyright 2026 Daniel Cederberg and William Zhang + * + * This file is part of the DNLP-differentiation-engine project. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include "affine.h" +#include +#include + +/* + * vstack(args) = transpose(hstack(transpose(args[0]), + * transpose(args[1]), + * ...)) + * + * All args must share the same d2 (number of columns). + */ +expr *new_vstack(expr **args, int n_args, int n_vars) +{ + assert(n_args > 0); + for (int i = 1; i < n_args; i++) + { + assert(args[i]->d2 == args[0]->d2); + } + + expr **transposed = (expr **) malloc(n_args * sizeof(expr *)); + for (int i = 0; i < n_args; i++) + { + transposed[i] = new_transpose(args[i]); + } + + expr *hstacked = new_hstack(transposed, n_args, n_vars); + free(transposed); + + return new_transpose(hstacked); +} diff --git a/tests/all_tests.c b/tests/all_tests.c index 8f839f49..414b80fa 100644 --- a/tests/all_tests.c +++ b/tests/all_tests.c @@ -7,6 +7,7 @@ #include "forward_pass/affine/test_add.h" #include "forward_pass/affine/test_broadcast.h" #include "forward_pass/affine/test_hstack.h" +#include "forward_pass/affine/test_vstack.h" #include "forward_pass/affine/test_linear_op.h" #include "forward_pass/affine/test_neg.h" #include "forward_pass/affine/test_promote.h" @@ -26,6 +27,7 @@ #include "jacobian_tests/test_const_vector_mult.h" #include "jacobian_tests/test_elementwise_mult.h" #include "jacobian_tests/test_hstack.h" +#include "jacobian_tests/test_vstack.h" #include "jacobian_tests/test_index.h" #include "jacobian_tests/test_left_matmul.h" #include "jacobian_tests/test_log.h" @@ -64,6 +66,7 @@ #include "wsum_hess/test_const_scalar_mult.h" #include "wsum_hess/test_const_vector_mult.h" #include "wsum_hess/test_hstack.h" +#include "wsum_hess/test_vstack.h" #include "wsum_hess/test_index.h" #include "wsum_hess/test_left_matmul.h" #include "wsum_hess/test_matmul.h" @@ -109,6 +112,8 @@ int main(void) mu_run_test(test_sum_axis_1, tests_run); mu_run_test(test_hstack_forward_vectors, tests_run); mu_run_test(test_hstack_forward_matrix, tests_run); + mu_run_test(test_vstack_forward_vectors, tests_run); + mu_run_test(test_vstack_forward_matrix, tests_run); mu_run_test(test_broadcast_row, tests_run); mu_run_test(test_broadcast_col, tests_run); mu_run_test(test_broadcast_matrix, tests_run); @@ -158,6 +163,8 @@ int main(void) mu_run_test(test_jacobian_sum_log_axis_1, tests_run); mu_run_test(test_jacobian_hstack_vectors, tests_run); mu_run_test(test_jacobian_hstack_matrix, tests_run); + mu_run_test(test_jacobian_vstack_vectors, tests_run); + mu_run_test(test_jacobian_vstack_matrix, tests_run); mu_run_test(test_index_forward_simple, tests_run); mu_run_test(test_index_forward_repeated, tests_run); mu_run_test(test_index_jacobian_of_variable, tests_run); @@ -218,6 +225,8 @@ int main(void) mu_run_test(test_wsum_hess_rel_entr_scalar_vector, tests_run); mu_run_test(test_wsum_hess_hstack, tests_run); mu_run_test(test_wsum_hess_hstack_matrix, tests_run); + mu_run_test(test_wsum_hess_vstack_vectors, tests_run); + mu_run_test(test_wsum_hess_vstack_matrix, tests_run); mu_run_test(test_wsum_hess_index_log, tests_run); mu_run_test(test_wsum_hess_index_repeated, tests_run); mu_run_test(test_wsum_hess_sum_index_log, tests_run); diff --git a/tests/forward_pass/affine/test_vstack.h b/tests/forward_pass/affine/test_vstack.h new file mode 100644 index 00000000..bd012408 --- /dev/null +++ b/tests/forward_pass/affine/test_vstack.h @@ -0,0 +1,87 @@ +#include +#include +#include + +#include "affine.h" +#include "elementwise_univariate.h" +#include "expr.h" +#include "minunit.h" +#include "test_helpers.h" + +const char *test_vstack_forward_vectors(void) +{ + /* vstack([log(x), exp(x)]) where x is (3,1) + * Output is (6,1): [log(1), log(2), log(3), exp(1), exp(2), exp(3)] + * For d2=1, vstack == hstack (no interleaving needed) + */ + double u[3] = {1.0, 2.0, 3.0}; + expr *x = new_variable(3, 1, 0, 3); + + expr *log_x = new_log(x); + expr *exp_x = new_exp(x); + + expr *args[2] = {log_x, exp_x}; + expr *stack = new_vstack(args, 2, 3); + + mu_assert("vstack vectors: wrong d1", stack->d1 == 6); + mu_assert("vstack vectors: wrong d2", stack->d2 == 1); + + stack->forward(stack, u); + + double expected[6] = {log(1.0), log(2.0), log(3.0), + exp(1.0), exp(2.0), exp(3.0)}; + + mu_assert("vstack forward vectors failed", + cmp_double_array(stack->value, expected, 6)); + + free_expr(stack); + return 0; +} + +const char *test_vstack_forward_matrix(void) +{ + /* vstack([log(x), exp(y)]) where x is (2,3), y is (1,3) + * x stored column-wise: [1, 2, 3, 4, 5, 6] + * x = [1 3 5] + * [2 4 6] + * y stored column-wise: [7, 8, 9] + * y = [7 8 9] + * + * Output is (3,3), stored column-wise: + * result = [log(1) log(3) log(5)] + * [log(2) log(4) log(6)] + * [exp(7) exp(8) exp(9)] + * + * Flat column-wise: + * [log(1), log(2), exp(7), log(3), log(4), exp(8), + * log(5), log(6), exp(9)] + */ + double ux[6] = {1.0, 2.0, 3.0, 4.0, 5.0, 6.0}; + double uy[3] = {7.0, 8.0, 9.0}; + double u[9] = {1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0}; + (void) ux; + (void) uy; + + expr *x = new_variable(2, 3, 0, 9); + expr *y = new_variable(1, 3, 6, 9); + + expr *log_x = new_log(x); + expr *exp_y = new_exp(y); + + expr *args[2] = {log_x, exp_y}; + expr *stack = new_vstack(args, 2, 9); + + mu_assert("vstack matrix: wrong d1", stack->d1 == 3); + mu_assert("vstack matrix: wrong d2", stack->d2 == 3); + + stack->forward(stack, u); + + double expected[9] = {log(1.0), log(2.0), exp(7.0), log(3.0), log(4.0), + exp(8.0), log(5.0), log(6.0), exp(9.0)}; + + mu_assert("vstack forward matrix failed", + cmp_double_array(stack->value, expected, 9)); + + free_expr(stack); + return 0; +} diff --git a/tests/jacobian_tests/test_vstack.h b/tests/jacobian_tests/test_vstack.h new file mode 100644 index 00000000..59dea4eb --- /dev/null +++ b/tests/jacobian_tests/test_vstack.h @@ -0,0 +1,99 @@ +#include +#include + +#include "affine.h" +#include "elementwise_univariate.h" +#include "expr.h" +#include "minunit.h" +#include "test_helpers.h" + +const char *test_jacobian_vstack_vectors(void) +{ + /* vstack([log(x), exp(x)]) where x is (3,1) + * Output (6,1) flat: [log(1), log(2), log(3), exp(1), exp(2), exp(3)] + * + * Jacobian is 6x3: + * row 0: d(log(x0))/dx0 = 1/1 at col 0 + * row 1: d(log(x1))/dx1 = 1/2 at col 1 + * row 2: d(log(x2))/dx2 = 1/3 at col 2 + * row 3: d(exp(x0))/dx0 = e^1 at col 0 + * row 4: d(exp(x1))/dx1 = e^2 at col 1 + * row 5: d(exp(x2))/dx2 = e^3 at col 2 + */ + double u[3] = {1.0, 2.0, 3.0}; + expr *x = new_variable(3, 1, 0, 3); + + expr *log_x = new_log(x); + expr *exp_x = new_exp(x); + + expr *args[2] = {log_x, exp_x}; + expr *stack = new_vstack(args, 2, 3); + + stack->forward(stack, u); + stack->jacobian_init(stack); + stack->eval_jacobian(stack); + + double expected_x[6] = {1.0, 0.5, 1.0 / 3.0, exp(1.0), exp(2.0), exp(3.0)}; + int expected_i[6] = {0, 1, 2, 0, 1, 2}; + int expected_p[7] = {0, 1, 2, 3, 4, 5, 6}; + + mu_assert("vstack jac vectors: vals", + cmp_double_array(stack->jacobian->x, expected_x, 6)); + mu_assert("vstack jac vectors: cols", + cmp_int_array(stack->jacobian->i, expected_i, 6)); + mu_assert("vstack jac vectors: rows", + cmp_int_array(stack->jacobian->p, expected_p, 7)); + + free_expr(stack); + return 0; +} + +const char *test_jacobian_vstack_matrix(void) +{ + /* vstack([log(x), exp(y)]) where x is (2,3), y is (1,3) + * x at var_id 0 (6 vars), y at var_id 6 (3 vars), total 9 vars + * + * Output (3,3) flat column-wise: + * [log(x0), log(x1), exp(y0), log(x2), log(x3), exp(y1), + * log(x4), log(x5), exp(y2)] + * + * Jacobian is 9x9 sparse: + * row 0 (log(x0)): 1/x0 at col 0 + * row 1 (log(x1)): 1/x1 at col 1 + * row 2 (exp(y0)): e^y0 at col 6 + * row 3 (log(x2)): 1/x2 at col 2 + * row 4 (log(x3)): 1/x3 at col 3 + * row 5 (exp(y1)): e^y1 at col 7 + * row 6 (log(x4)): 1/x4 at col 4 + * row 7 (log(x5)): 1/x5 at col 5 + * row 8 (exp(y2)): e^y2 at col 8 + */ + double u[9] = {1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0}; + expr *x = new_variable(2, 3, 0, 9); + expr *y = new_variable(1, 3, 6, 9); + + expr *log_x = new_log(x); + expr *exp_y = new_exp(y); + + expr *args[2] = {log_x, exp_y}; + expr *stack = new_vstack(args, 2, 9); + + stack->forward(stack, u); + stack->jacobian_init(stack); + stack->eval_jacobian(stack); + + double expected_x[9] = {1.0, 0.5, exp(7.0), 1.0 / 3.0, 0.25, + exp(8.0), 0.2, 1.0 / 6.0, exp(9.0)}; + int expected_i[9] = {0, 1, 6, 2, 3, 7, 4, 5, 8}; + int expected_p[10] = {0, 1, 2, 3, 4, 5, 6, 7, 8, 9}; + + mu_assert("vstack jac matrix: vals", + cmp_double_array(stack->jacobian->x, expected_x, 9)); + mu_assert("vstack jac matrix: cols", + cmp_int_array(stack->jacobian->i, expected_i, 9)); + mu_assert("vstack jac matrix: rows", + cmp_int_array(stack->jacobian->p, expected_p, 10)); + + free_expr(stack); + return 0; +} diff --git a/tests/wsum_hess/test_vstack.h b/tests/wsum_hess/test_vstack.h new file mode 100644 index 00000000..e3ce6aec --- /dev/null +++ b/tests/wsum_hess/test_vstack.h @@ -0,0 +1,113 @@ +#include "affine.h" +#include "elementwise_univariate.h" +#include "expr.h" +#include "minunit.h" +#include "test_helpers.h" +#include +#include + +const char *test_wsum_hess_vstack_vectors(void) +{ + /* vstack([log(x), exp(x)]) where x is (3,1) + * Output (6,1): [log(x0), log(x1), log(x2), + * exp(x0), exp(x1), exp(x2)] + * w = [1, 2, 3, 4, 5, 6] + * + * Hessian (3x3 diagonal): + * (0,0): w[0]*(-1/x0^2) + w[3]*exp(x0) = -1 + 4*e + * (1,1): w[1]*(-1/x1^2) + w[4]*exp(x1) = -2/4 + 5*e^2 + * (2,2): w[2]*(-1/x2^2) + w[5]*exp(x2) = -3/9 + 6*e^3 + */ + double u[3] = {1.0, 2.0, 3.0}; + double w[6] = {1.0, 2.0, 3.0, 4.0, 5.0, 6.0}; + + expr *x = new_variable(3, 1, 0, 3); + expr *log_x = new_log(x); + expr *exp_x = new_exp(x); + + expr *args[2] = {log_x, exp_x}; + expr *stack = new_vstack(args, 2, 3); + + stack->forward(stack, u); + stack->jacobian_init(stack); + stack->wsum_hess_init(stack); + stack->eval_wsum_hess(stack, w); + + double expected_x[3] = {-1.0 + 4.0 * exp(1.0), -0.5 + 5.0 * exp(2.0), + -1.0 / 3.0 + 6.0 * exp(3.0)}; + int expected_p[4] = {0, 1, 2, 3}; + int expected_i[3] = {0, 1, 2}; + + mu_assert("vstack hess vectors: vals", + cmp_double_array(stack->wsum_hess->x, expected_x, 3)); + mu_assert("vstack hess vectors: rows", + cmp_int_array(stack->wsum_hess->p, expected_p, 4)); + mu_assert("vstack hess vectors: cols", + cmp_int_array(stack->wsum_hess->i, expected_i, 3)); + + free_expr(stack); + return 0; +} + +const char *test_wsum_hess_vstack_matrix(void) +{ + /* vstack([log(x), exp(y)]) where x is (2,3), y is (1,3) + * x at var_id 0 (6 vars), y at var_id 6 (3 vars) + * + * Output (3,3) flat column-wise: + * [log(x0), log(x1), exp(y0), log(x2), log(x3), exp(y1), + * log(x4), log(x5), exp(y2)] + * w = [1, 2, 3, 4, 5, 6, 7, 8, 9] + * + * Hessian (9x9 diagonal): + * x vars (0-5): w[k]*(-1/xk^2) where k maps output→var + * (0,0): w[0]*(-1/1^2) = -1 + * (1,1): w[1]*(-1/2^2) = -0.5 + * (2,2): w[3]*(-1/3^2) = -4/9 + * (3,3): w[4]*(-1/4^2) = -5/16 + * (4,4): w[6]*(-1/5^2) = -7/25 + * (5,5): w[7]*(-1/6^2) = -8/36 + * y vars (6-8): w[k]*exp(yk) where k maps output→var + * (6,6): w[2]*exp(7) = 3*exp(7) + * (7,7): w[5]*exp(8) = 6*exp(8) + * (8,8): w[8]*exp(9) = 9*exp(9) + */ + double u[9] = {1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0}; + double w[9] = {1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0}; + + expr *x = new_variable(2, 3, 0, 9); + expr *y = new_variable(1, 3, 6, 9); + + expr *log_x = new_log(x); + expr *exp_y = new_exp(y); + + expr *args[2] = {log_x, exp_y}; + expr *stack = new_vstack(args, 2, 9); + + stack->forward(stack, u); + stack->jacobian_init(stack); + stack->wsum_hess_init(stack); + stack->eval_wsum_hess(stack, w); + + double expected_x[9] = {-1.0, /* x0: w[0]*(-1/1) */ + -0.5, /* x1: w[1]*(-1/4) */ + -4.0 / 9.0, /* x2: w[3]*(-1/9) */ + -5.0 / 16.0, /* x3: w[4]*(-1/16) */ + -7.0 / 25.0, /* x4: w[6]*(-1/25) */ + -8.0 / 36.0, /* x5: w[7]*(-1/36) */ + 3.0 * exp(7.0), /* y0: w[2]*exp(7) */ + 6.0 * exp(8.0), /* y1: w[5]*exp(8) */ + 9.0 * exp(9.0)}; /* y2: w[8]*exp(9) */ + int expected_p[10] = {0, 1, 2, 3, 4, 5, 6, 7, 8, 9}; + int expected_i[9] = {0, 1, 2, 3, 4, 5, 6, 7, 8}; + + mu_assert("vstack hess matrix: vals", + cmp_double_array(stack->wsum_hess->x, expected_x, 9)); + mu_assert("vstack hess matrix: rows", + cmp_int_array(stack->wsum_hess->p, expected_p, 10)); + mu_assert("vstack hess matrix: cols", + cmp_int_array(stack->wsum_hess->i, expected_i, 9)); + + free_expr(stack); + return 0; +} From 045f86fc9dee1e80bf844cf0c1f7659567935ef6 Mon Sep 17 00:00:00 2001 From: William Zijie Zhang Date: Thu, 26 Mar 2026 10:37:46 -0400 Subject: [PATCH 3/3] Remove unused ux and uy variables from vstack forward test Co-Authored-By: Claude Opus 4.6 (1M context) --- tests/forward_pass/affine/test_vstack.h | 4 ---- 1 file changed, 4 deletions(-) diff --git a/tests/forward_pass/affine/test_vstack.h b/tests/forward_pass/affine/test_vstack.h index bd012408..fefb80f4 100644 --- a/tests/forward_pass/affine/test_vstack.h +++ b/tests/forward_pass/affine/test_vstack.h @@ -56,11 +56,7 @@ const char *test_vstack_forward_matrix(void) * [log(1), log(2), exp(7), log(3), log(4), exp(8), * log(5), log(6), exp(9)] */ - double ux[6] = {1.0, 2.0, 3.0, 4.0, 5.0, 6.0}; - double uy[3] = {7.0, 8.0, 9.0}; double u[9] = {1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0}; - (void) ux; - (void) uy; expr *x = new_variable(2, 3, 0, 9); expr *y = new_variable(1, 3, 6, 9);