From a5dfe1b24c731ee6bc7424a1c1357119df9d1e01 Mon Sep 17 00:00:00 2001 From: William Zijie Zhang Date: Wed, 25 Mar 2026 18:53:52 -0400 Subject: [PATCH] Add upper_tri and diag_mat affine atoms Both are thin wrappers that compute flat column-major index arrays and delegate to new_index. diag_mat extracts the diagonal of a square matrix into a vector. upper_tri extracts strict upper triangular elements (excluding diagonal). Also removes a duplicate new_diag_vec declaration in affine.h. Co-Authored-By: Claude Opus 4.6 (1M context) --- include/affine.h | 3 +- src/affine/diag_mat.c | 42 ++++++++++++++ src/affine/upper_tri.c | 54 ++++++++++++++++++ tests/all_tests.c | 14 +++++ tests/forward_pass/affine/test_diag_mat.h | 31 +++++++++++ tests/forward_pass/affine/test_upper_tri.h | 32 +++++++++++ tests/jacobian_tests/test_diag_mat.h | 63 +++++++++++++++++++++ tests/jacobian_tests/test_upper_tri.h | 65 ++++++++++++++++++++++ tests/wsum_hess/test_diag_mat.h | 48 ++++++++++++++++ tests/wsum_hess/test_upper_tri.h | 56 +++++++++++++++++++ 10 files changed, 407 insertions(+), 1 deletion(-) create mode 100644 src/affine/diag_mat.c create mode 100644 src/affine/upper_tri.c create mode 100644 tests/forward_pass/affine/test_diag_mat.h create mode 100644 tests/forward_pass/affine/test_upper_tri.h create mode 100644 tests/jacobian_tests/test_diag_mat.h create mode 100644 tests/jacobian_tests/test_upper_tri.h create mode 100644 tests/wsum_hess/test_diag_mat.h create mode 100644 tests/wsum_hess/test_upper_tri.h diff --git a/include/affine.h b/include/affine.h index ada58f1..e68da01 100644 --- a/include/affine.h +++ b/include/affine.h @@ -39,7 +39,8 @@ expr *new_index(expr *child, int d1, int d2, const int *indices, int n_idxs); expr *new_reshape(expr *child, int d1, int d2); expr *new_broadcast(expr *child, int target_d1, int target_d2); expr *new_diag_vec(expr *child); +expr *new_diag_mat(expr *child); +expr *new_upper_tri(expr *child); expr *new_transpose(expr *child); -expr *new_diag_vec(expr *child); #endif /* AFFINE_H */ diff --git a/src/affine/diag_mat.c b/src/affine/diag_mat.c new file mode 100644 index 0000000..f16cfb3 --- /dev/null +++ b/src/affine/diag_mat.c @@ -0,0 +1,42 @@ +/* + * 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. + */ +// SPDX-License-Identifier: Apache-2.0 + +#include "affine.h" +#include +#include + +/* Extract diagonal from a square matrix into a column vector. + * For an (n, n) matrix in column-major order, diagonal element i + * is at flat index i * (n + 1). */ + +expr *new_diag_mat(expr *child) +{ + assert(child->d1 == child->d2); + int n = child->d1; + + int *indices = (int *) malloc((size_t) n * sizeof(int)); + for (int i = 0; i < n; i++) + { + indices[i] = i * (n + 1); + } + + expr *node = new_index(child, n, 1, indices, n); + free(indices); + return node; +} diff --git a/src/affine/upper_tri.c b/src/affine/upper_tri.c new file mode 100644 index 0000000..82ac948 --- /dev/null +++ b/src/affine/upper_tri.c @@ -0,0 +1,54 @@ +/* + * 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. + */ +// SPDX-License-Identifier: Apache-2.0 + +#include "affine.h" +#include +#include + +/* Extract strict upper triangular elements (excluding diagonal) + * from a square matrix, in column-major order. + * + * For an (n, n) matrix, element (i, j) with i < j is at flat + * index j * n + i. Output has n * (n - 1) / 2 elements. */ + +expr *new_upper_tri(expr *child) +{ + assert(child->d1 == child->d2); + int n = child->d1; + int n_elems = n * (n - 1) / 2; + + int *indices = NULL; + if (n_elems > 0) + { + indices = (int *) malloc((size_t) n_elems * sizeof(int)); + int k = 0; + for (int j = 0; j < n; j++) + { + for (int i = 0; i < j; i++) + { + indices[k++] = j * n + i; + } + } + assert(k == n_elems); + } + + expr *node = new_index(child, n_elems, 1, indices, n_elems); + free(indices); + return node; +} diff --git a/tests/all_tests.c b/tests/all_tests.c index 8f839f4..8afad2a 100644 --- a/tests/all_tests.c +++ b/tests/all_tests.c @@ -6,11 +6,13 @@ #ifndef PROFILE_ONLY #include "forward_pass/affine/test_add.h" #include "forward_pass/affine/test_broadcast.h" +#include "forward_pass/affine/test_diag_mat.h" #include "forward_pass/affine/test_hstack.h" #include "forward_pass/affine/test_linear_op.h" #include "forward_pass/affine/test_neg.h" #include "forward_pass/affine/test_promote.h" #include "forward_pass/affine/test_sum.h" +#include "forward_pass/affine/test_upper_tri.h" #include "forward_pass/affine/test_variable_constant.h" #include "forward_pass/composite/test_composite.h" #include "forward_pass/elementwise/test_exp.h" @@ -24,6 +26,7 @@ #include "jacobian_tests/test_composite.h" #include "jacobian_tests/test_const_scalar_mult.h" #include "jacobian_tests/test_const_vector_mult.h" +#include "jacobian_tests/test_diag_mat.h" #include "jacobian_tests/test_elementwise_mult.h" #include "jacobian_tests/test_hstack.h" #include "jacobian_tests/test_index.h" @@ -44,6 +47,7 @@ #include "jacobian_tests/test_sum.h" #include "jacobian_tests/test_trace.h" #include "jacobian_tests/test_transpose.h" +#include "jacobian_tests/test_upper_tri.h" #include "problem/test_problem.h" #include "utils/test_cblas.h" #include "utils/test_coo_matrix.h" @@ -63,6 +67,7 @@ #include "wsum_hess/test_broadcast.h" #include "wsum_hess/test_const_scalar_mult.h" #include "wsum_hess/test_const_vector_mult.h" +#include "wsum_hess/test_diag_mat.h" #include "wsum_hess/test_hstack.h" #include "wsum_hess/test_index.h" #include "wsum_hess/test_left_matmul.h" @@ -80,6 +85,7 @@ #include "wsum_hess/test_sum.h" #include "wsum_hess/test_trace.h" #include "wsum_hess/test_transpose.h" +#include "wsum_hess/test_upper_tri.h" #endif /* PROFILE_ONLY */ #ifdef PROFILE_ONLY @@ -116,6 +122,8 @@ int main(void) mu_run_test(test_forward_prod_axis_one, tests_run); mu_run_test(test_matmul, tests_run); mu_run_test(test_left_matmul_dense, tests_run); + mu_run_test(test_diag_mat_forward, tests_run); + mu_run_test(test_upper_tri_forward, tests_run); printf("\n--- Jacobian Tests ---\n"); mu_run_test(test_neg_jacobian, tests_run); @@ -181,6 +189,10 @@ int main(void) mu_run_test(test_jacobian_right_matmul_log_vector, tests_run); mu_run_test(test_jacobian_matmul, tests_run); mu_run_test(test_jacobian_transpose, tests_run); + mu_run_test(test_diag_mat_jacobian_variable, tests_run); + mu_run_test(test_diag_mat_jacobian_of_log, tests_run); + mu_run_test(test_upper_tri_jacobian_variable, tests_run); + mu_run_test(test_upper_tri_jacobian_of_log, tests_run); printf("\n--- Weighted Sum of Hessian Tests ---\n"); mu_run_test(test_wsum_hess_log, tests_run); @@ -246,6 +258,8 @@ int main(void) mu_run_test(test_wsum_hess_trace_log_variable, tests_run); mu_run_test(test_wsum_hess_trace_composite, tests_run); mu_run_test(test_wsum_hess_transpose, tests_run); + mu_run_test(test_wsum_hess_diag_mat_log, tests_run); + mu_run_test(test_wsum_hess_upper_tri_log, tests_run); printf("\n--- Utility Tests ---\n"); mu_run_test(test_cblas_ddot, tests_run); diff --git a/tests/forward_pass/affine/test_diag_mat.h b/tests/forward_pass/affine/test_diag_mat.h new file mode 100644 index 0000000..3815861 --- /dev/null +++ b/tests/forward_pass/affine/test_diag_mat.h @@ -0,0 +1,31 @@ +// SPDX-License-Identifier: Apache-2.0 + +#include + +#include "affine.h" +#include "expr.h" +#include "minunit.h" +#include "test_helpers.h" + +const char *test_diag_mat_forward(void) +{ + /* 3x3 matrix variable (column-major): [1,2,3,4,5,6,7,8,9] + * Matrix: 1 4 7 + * 2 5 8 + * 3 6 9 + * Diagonal: (0,0)=1, (1,1)=5, (2,2)=9 */ + double u[9] = {1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0}; + expr *var = new_variable(3, 3, 0, 9); + expr *dm = new_diag_mat(var); + + mu_assert("diag_mat d1", dm->d1 == 3); + mu_assert("diag_mat d2", dm->d2 == 1); + + dm->forward(dm, u); + + double expected[3] = {1.0, 5.0, 9.0}; + mu_assert("diag_mat forward", cmp_double_array(dm->value, expected, 3)); + + free_expr(dm); + return 0; +} diff --git a/tests/forward_pass/affine/test_upper_tri.h b/tests/forward_pass/affine/test_upper_tri.h new file mode 100644 index 0000000..9663540 --- /dev/null +++ b/tests/forward_pass/affine/test_upper_tri.h @@ -0,0 +1,32 @@ +// SPDX-License-Identifier: Apache-2.0 + +#include + +#include "affine.h" +#include "expr.h" +#include "minunit.h" +#include "test_helpers.h" + +const char *test_upper_tri_forward(void) +{ + /* 3x3 matrix variable (column-major): [1,2,3,4,5,6,7,8,9] + * Matrix: 1 4 7 + * 2 5 8 + * 3 6 9 + * Upper tri (i < j): (0,1)=4, (0,2)=7, (1,2)=8 + * Flat indices: 3, 6, 7 */ + double u[9] = {1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0}; + expr *var = new_variable(3, 3, 0, 9); + expr *ut = new_upper_tri(var); + + mu_assert("upper_tri d1", ut->d1 == 3); + mu_assert("upper_tri d2", ut->d2 == 1); + + ut->forward(ut, u); + + double expected[3] = {4.0, 7.0, 8.0}; + mu_assert("upper_tri forward", cmp_double_array(ut->value, expected, 3)); + + free_expr(ut); + return 0; +} diff --git a/tests/jacobian_tests/test_diag_mat.h b/tests/jacobian_tests/test_diag_mat.h new file mode 100644 index 0000000..58143a7 --- /dev/null +++ b/tests/jacobian_tests/test_diag_mat.h @@ -0,0 +1,63 @@ +// SPDX-License-Identifier: Apache-2.0 + +#include + +#include "affine.h" +#include "elementwise_univariate.h" +#include "expr.h" +#include "minunit.h" +#include "test_helpers.h" + +const char *test_diag_mat_jacobian_variable(void) +{ + /* diag_mat of a 2x2 variable (4 vars total) + * Diagonal indices in column-major: [0, 3] + * Jacobian is 2x4 CSR: row 0 has col 0, row 1 has col 3 */ + double u[4] = {1.0, 2.0, 3.0, 4.0}; + expr *var = new_variable(2, 2, 0, 4); + expr *dm = new_diag_mat(var); + + dm->forward(dm, u); + dm->jacobian_init(dm); + dm->eval_jacobian(dm); + + double expected_x[2] = {1.0, 1.0}; + int expected_p[3] = {0, 1, 2}; + int expected_i[2] = {0, 3}; + + mu_assert("diag_mat jac vals", cmp_double_array(dm->jacobian->x, expected_x, 2)); + mu_assert("diag_mat jac p", cmp_int_array(dm->jacobian->p, expected_p, 3)); + mu_assert("diag_mat jac i", cmp_int_array(dm->jacobian->i, expected_i, 2)); + + free_expr(dm); + return 0; +} + +const char *test_diag_mat_jacobian_of_log(void) +{ + /* diag_mat(log(X)) where X is 2x2 variable + * X = [[1, 3], [2, 4]] (column-major: [1, 2, 3, 4]) + * Diagonal: x[0]=1, x[3]=4 + * d/dx log at diagonal positions: + * Row 0: 1/1 = 1.0 at col 0 + * Row 1: 1/4 = 0.25 at col 3 */ + double u[4] = {1.0, 2.0, 3.0, 4.0}; + expr *var = new_variable(2, 2, 0, 4); + expr *log_node = new_log(var); + expr *dm = new_diag_mat(log_node); + + dm->forward(dm, u); + dm->jacobian_init(dm); + dm->eval_jacobian(dm); + + double expected_x[2] = {1.0, 0.25}; + int expected_i[2] = {0, 3}; + + mu_assert("diag_mat log jac vals", + cmp_double_array(dm->jacobian->x, expected_x, 2)); + mu_assert("diag_mat log jac cols", + cmp_int_array(dm->jacobian->i, expected_i, 2)); + + free_expr(dm); + return 0; +} diff --git a/tests/jacobian_tests/test_upper_tri.h b/tests/jacobian_tests/test_upper_tri.h new file mode 100644 index 0000000..5f8d7e6 --- /dev/null +++ b/tests/jacobian_tests/test_upper_tri.h @@ -0,0 +1,65 @@ +// SPDX-License-Identifier: Apache-2.0 + +#include + +#include "affine.h" +#include "elementwise_univariate.h" +#include "expr.h" +#include "minunit.h" +#include "test_helpers.h" + +const char *test_upper_tri_jacobian_variable(void) +{ + /* upper_tri of a 3x3 variable (9 vars total) + * Upper tri flat indices: [3, 6, 7] + * Jacobian is 3x9 CSR: row k has col indices[k] */ + double u[9] = {1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0}; + expr *var = new_variable(3, 3, 0, 9); + expr *ut = new_upper_tri(var); + + ut->forward(ut, u); + ut->jacobian_init(ut); + ut->eval_jacobian(ut); + + double expected_x[3] = {1.0, 1.0, 1.0}; + int expected_p[4] = {0, 1, 2, 3}; + int expected_i[3] = {3, 6, 7}; + + mu_assert("upper_tri jac vals", + cmp_double_array(ut->jacobian->x, expected_x, 3)); + mu_assert("upper_tri jac p", cmp_int_array(ut->jacobian->p, expected_p, 4)); + mu_assert("upper_tri jac i", cmp_int_array(ut->jacobian->i, expected_i, 3)); + + free_expr(ut); + return 0; +} + +const char *test_upper_tri_jacobian_of_log(void) +{ + /* upper_tri(log(X)) where X is 3x3 variable + * Upper tri flat indices: [3, 6, 7] + * X values at those positions: x[3]=4, x[6]=7, x[7]=8 + * d/dx log at those positions: + * Row 0: 1/4 = 0.25 at col 3 + * Row 1: 1/7 at col 6 + * Row 2: 1/8 = 0.125 at col 7 */ + double u[9] = {1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0}; + expr *var = new_variable(3, 3, 0, 9); + expr *log_node = new_log(var); + expr *ut = new_upper_tri(log_node); + + ut->forward(ut, u); + ut->jacobian_init(ut); + ut->eval_jacobian(ut); + + double expected_x[3] = {0.25, 1.0 / 7.0, 0.125}; + int expected_i[3] = {3, 6, 7}; + + mu_assert("upper_tri log jac vals", + cmp_double_array(ut->jacobian->x, expected_x, 3)); + mu_assert("upper_tri log jac cols", + cmp_int_array(ut->jacobian->i, expected_i, 3)); + + free_expr(ut); + return 0; +} diff --git a/tests/wsum_hess/test_diag_mat.h b/tests/wsum_hess/test_diag_mat.h new file mode 100644 index 0000000..ce6a0b3 --- /dev/null +++ b/tests/wsum_hess/test_diag_mat.h @@ -0,0 +1,48 @@ +// SPDX-License-Identifier: Apache-2.0 + +#include +#include +#include +#include + +#include "affine.h" +#include "elementwise_univariate.h" +#include "expr.h" +#include "minunit.h" +#include "test_helpers.h" + +const char *test_wsum_hess_diag_mat_log(void) +{ + /* diag_mat(log(X)) where X is 2x2, w = [1, 1] + * X = [[1, 3], [2, 4]] (column-major: [1, 2, 3, 4]) + * Diagonal positions: flat indices 0 and 3 + * Hessian of log is diag(-1/x^2) + * Weights scatter: parent_w = [1, 0, 0, 1] + * H[0,0] = -1 * 1/1 = -1.0 + * H[1,1] = 0 (no weight) + * H[2,2] = 0 (no weight) + * H[3,3] = -1 * 1/16 = -0.0625 */ + double u[4] = {1.0, 2.0, 3.0, 4.0}; + double w[2] = {1.0, 1.0}; + + expr *var = new_variable(2, 2, 0, 4); + expr *log_node = new_log(var); + expr *dm = new_diag_mat(log_node); + + dm->forward(dm, u); + dm->jacobian_init(dm); + dm->wsum_hess_init(dm); + dm->eval_wsum_hess(dm, w); + + double expected_x[4] = {-1.0, 0.0, 0.0, -0.0625}; + int expected_p[5] = {0, 1, 2, 3, 4}; + int expected_i[4] = {0, 1, 2, 3}; + + mu_assert("diag_mat log hess vals", + cmp_double_array(dm->wsum_hess->x, expected_x, 4)); + mu_assert("diag_mat log hess p", cmp_int_array(dm->wsum_hess->p, expected_p, 5)); + mu_assert("diag_mat log hess i", cmp_int_array(dm->wsum_hess->i, expected_i, 4)); + + free_expr(dm); + return 0; +} diff --git a/tests/wsum_hess/test_upper_tri.h b/tests/wsum_hess/test_upper_tri.h new file mode 100644 index 0000000..ee6ff0e --- /dev/null +++ b/tests/wsum_hess/test_upper_tri.h @@ -0,0 +1,56 @@ +// SPDX-License-Identifier: Apache-2.0 + +#include +#include +#include +#include + +#include "affine.h" +#include "elementwise_univariate.h" +#include "expr.h" +#include "minunit.h" +#include "test_helpers.h" + +const char *test_wsum_hess_upper_tri_log(void) +{ + /* upper_tri(log(X)) where X is 3x3, w = [1, 1, 1] + * X (column-major): [1, 2, 3, 4, 5, 6, 7, 8, 9] + * Upper tri flat indices: [3, 6, 7] + * Hessian of log is diag(-1/x^2) + * Weights scatter: parent_w[3]=1, parent_w[6]=1, parent_w[7]=1 + * All other parent_w entries = 0 + * + * H[k,k] = parent_w[k] * (-1/x[k]^2): + * H[0,0] = 0, H[1,1] = 0, H[2,2] = 0 + * H[3,3] = -1/16 = -0.0625 + * H[4,4] = 0, H[5,5] = 0 + * H[6,6] = -1/49 + * H[7,7] = -1/64 = -0.015625 + * H[8,8] = 0 */ + double u[9] = {1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0}; + double w[3] = {1.0, 1.0, 1.0}; + + expr *var = new_variable(3, 3, 0, 9); + expr *log_node = new_log(var); + expr *ut = new_upper_tri(log_node); + + ut->forward(ut, u); + ut->jacobian_init(ut); + ut->wsum_hess_init(ut); + ut->eval_wsum_hess(ut, w); + + double expected_x[9] = {0.0, 0.0, 0.0, -0.0625, 0.0, + 0.0, -1.0 / 49.0, -0.015625, 0.0}; + 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("upper_tri log hess vals", + cmp_double_array(ut->wsum_hess->x, expected_x, 9)); + mu_assert("upper_tri log hess p", + cmp_int_array(ut->wsum_hess->p, expected_p, 10)); + mu_assert("upper_tri log hess i", + cmp_int_array(ut->wsum_hess->i, expected_i, 9)); + + free_expr(ut); + return 0; +}