[Ready for review] Add vstack atom through transpose(hstack(transpose[args]))

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);
 

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 <assert.h>
+#include <stdlib.h>
+
+/*
+ * 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);
+}

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);

tests/forward_pass/affine/test_vstack.h

@@ -0,0 +1,83 @@
+#include <math.h>
+#include <stdio.h>
+#include <stdlib.h>
+
+#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;
+}

tests/jacobian_tests/test_vstack.h

@@ -0,0 +1,99 @@
+#include <math.h>
+#include <stdio.h>
+
+#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;
+}