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..fefb80f4 --- /dev/null +++ b/tests/forward_pass/affine/test_vstack.h @@ -0,0 +1,83 @@ +#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 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); + + 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; +}