[WIP] Add upper_tri and diag_mat affine atoms

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 */

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

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

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

tests/forward_pass/affine/test_diag_mat.h

@@ -0,0 +1,31 @@
+// SPDX-License-Identifier: Apache-2.0
+
+#include <stdio.h>
+
+#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;
+}

tests/forward_pass/affine/test_upper_tri.h

@@ -0,0 +1,32 @@
+// SPDX-License-Identifier: Apache-2.0
+
+#include <stdio.h>
+
+#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;
+}

tests/jacobian_tests/test_diag_mat.h

@@ -0,0 +1,63 @@
+// SPDX-License-Identifier: Apache-2.0
+
+#include <stdio.h>
+
+#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;
+}

tests/jacobian_tests/test_upper_tri.h

@@ -0,0 +1,65 @@
+// SPDX-License-Identifier: Apache-2.0
+
+#include <stdio.h>
+
+#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;
+}