From 02aeea999ade87f289c7e7df843c958927644cb1 Mon Sep 17 00:00:00 2001 From: Daniel Date: Thu, 15 Jan 2026 10:56:21 -0800 Subject: [PATCH 1/8] going through dimensions of atoms --- TODO.md | 1 + src/affine/broadcast.c | 10 ---------- src/bivariate/quad_over_lin.c | 8 ++++++-- src/other/quad_form.c | 21 ++++++++++++--------- tests/jacobian_tests/test_index.h | 2 +- 5 files changed, 20 insertions(+), 22 deletions(-) diff --git a/TODO.md b/TODO.md index d473508..ec3ed93 100644 --- a/TODO.md +++ b/TODO.md @@ -4,6 +4,7 @@ 6. AX where X is a matrix. Can that happen? How is that canonicalized? Maybe it can't happen. 7. Must be able to compute jacobian and hessian of A @ phi(x), so linear operator needs other code! This requires new infrastructure, I think. 8. Shortcut hessians of affine stuff? +9. quad-over-lin with first argument matrix. Going through all atoms to see that sparsity pattern is computed in initialization of jacobian: 2. trace - not ok diff --git a/src/affine/broadcast.c b/src/affine/broadcast.c index 6739983..8831c9c 100644 --- a/src/affine/broadcast.c +++ b/src/affine/broadcast.c @@ -250,16 +250,6 @@ expr *new_broadcast(expr *child, int target_d1, int target_d2) { type = BROADCAST_SCALAR; } - else if (child->d1 == n && child->d2 == 1) - { - /* cvxpy convention requires this sometimes */ - type = BROADCAST_ROW; - } - else if (child->d1 == 1 && child->d2 == m) - { - /* cvxpy convention requires this sometimes */ - type = BROADCAST_COL; - } else { fprintf( diff --git a/src/bivariate/quad_over_lin.c b/src/bivariate/quad_over_lin.c index 41b8dc5..1e36e29 100644 --- a/src/bivariate/quad_over_lin.c +++ b/src/bivariate/quad_over_lin.c @@ -22,7 +22,7 @@ static void forward(expr *node, const double *u) /* local forward pass */ node->value[0] = 0.0; - for (int i = 0; i < x->d1; i++) + for (int i = 0; i < x->size; i++) { node->value[0] += x->value[i] * x->value[i]; } @@ -291,9 +291,13 @@ static void eval_wsum_hess(expr *node, const double *w) } } +/* TODO: the above implementations probably assumes that the first argument is a + vector but cvxpy supports a matrix for the first argument*/ + expr *new_quad_over_lin(expr *left, expr *right) { - expr *node = new_expr(left->d1, 1, left->n_vars); + assert((left->d2 == 1 && right->d2 == 1)); /* right must be scalar*/ + expr *node = new_expr(1, 1, left->n_vars); node->left = left; node->right = right; expr_retain(left); diff --git a/src/other/quad_form.c b/src/other/quad_form.c index 1fb68a3..da7fbae 100644 --- a/src/other/quad_form.c +++ b/src/other/quad_form.c @@ -170,25 +170,28 @@ static void eval_jacobian_old(expr *node) } */ +static void free_type_data(expr *node) +{ + quad_form_expr *qnode = (quad_form_expr *) node; + free_csr_matrix(qnode->Q); + qnode->Q = NULL; +} + expr *new_quad_form(expr *left, CSR_Matrix *Q) { + assert(left->d2 == 1); quad_form_expr *qnode = (quad_form_expr *) calloc(1, sizeof(quad_form_expr)); expr *node = &qnode->base; - /* Initialize base fields */ - assert(left->d2 == 1); - init_expr(node, left->d1, 1, left->n_vars, forward, jacobian_init, eval_jacobian, - NULL, NULL); - - /* Set left child */ + init_expr(node, 1, 1, left->n_vars, forward, jacobian_init, eval_jacobian, NULL, + free_type_data); node->left = left; expr_retain(left); + node->wsum_hess_init = wsum_hess_init; + node->eval_wsum_hess = eval_wsum_hess; /* Set type-specific field */ qnode->Q = new_csr_matrix(Q->m, Q->n, Q->nnz); copy_csr_matrix(Q, qnode->Q); - - node->wsum_hess_init = wsum_hess_init; - node->eval_wsum_hess = eval_wsum_hess; return node; } diff --git a/tests/jacobian_tests/test_index.h b/tests/jacobian_tests/test_index.h index 7f3c905..00b7cb2 100644 --- a/tests/jacobian_tests/test_index.h +++ b/tests/jacobian_tests/test_index.h @@ -111,7 +111,7 @@ const char *test_index_jacobian_repeated(void) mu_assert("index repeated jac vals", cmp_double_array(idx->jacobian->x, expected_x, 2)); mu_assert("index repeated row ptr", - cmp_int_array(idx->jacobian->p, expected_p, 4)); + cmp_int_array(idx->jacobian->p, expected_p, 3)); mu_assert("index repeated jac i", cmp_int_array(idx->jacobian->i, expected_i, 2)); From f20a725984c35c21406cb4558d2a935414f75281 Mon Sep 17 00:00:00 2001 From: Daniel Date: Fri, 16 Jan 2026 03:27:34 -0800 Subject: [PATCH 2/8] quad form, index dimension etc --- src/affine/broadcast.c | 1 + src/affine/index.c | 6 ++++-- src/affine/sum.c | 21 +++++++++++---------- src/other/quad_form.c | 2 +- 4 files changed, 17 insertions(+), 13 deletions(-) diff --git a/src/affine/broadcast.c b/src/affine/broadcast.c index 8831c9c..a26bd35 100644 --- a/src/affine/broadcast.c +++ b/src/affine/broadcast.c @@ -100,6 +100,7 @@ static void jacobian_init(expr *node) J->nnz += nnz_in_row; } } + J->p[node->size] = total_nnz; } else if (bcast->type == BROADCAST_COL) { diff --git a/src/affine/index.c b/src/affine/index.c index 81def9b..c8f6ade 100644 --- a/src/affine/index.c +++ b/src/affine/index.c @@ -148,8 +148,10 @@ expr *new_index(expr *child, const int *indices, int n_idxs) index_expr *idx = (index_expr *) calloc(1, sizeof(index_expr)); expr *node = &idx->base; - /* output shape is (n_idxs, 1) - flattened */ - init_expr(node, n_idxs, 1, child->n_vars, forward, jacobian_init, eval_jacobian, + /* to be consistent with CVXPY and NumPy we treat the result from + indexing as a row vector. + TODO: is this correct? What if the index in numpy/cvxpy returns a matrix? */ + init_expr(node, 1, n_idxs, child->n_vars, forward, jacobian_init, eval_jacobian, is_affine, free_type_data); node->left = child; diff --git a/src/affine/sum.c b/src/affine/sum.c index 9ab1f60..f30e0f3 100644 --- a/src/affine/sum.c +++ b/src/affine/sum.c @@ -45,7 +45,7 @@ static void forward(expr *node, const double *u) } else if (axis == 1) { - memset(node->value, 0, node->d1 * sizeof(double)); + memset(node->value, 0, node->size * sizeof(double)); /* sum columns together */ for (j = 0; j < x->d2; j++) @@ -69,7 +69,7 @@ static void jacobian_init(expr *node) x->jacobian_init(x); /* we never have to store more than the child's nnz */ - node->jacobian = new_csr_matrix(node->d1, node->n_vars, x->jacobian->nnz); + node->jacobian = new_csr_matrix(node->size, node->n_vars, x->jacobian->nnz); node->iwork = malloc(MAX(node->jacobian->n, x->jacobian->nnz) * sizeof(int)); snode->idx_map = malloc(x->jacobian->nnz * sizeof(int)); @@ -86,7 +86,7 @@ static void jacobian_init(expr *node) else if (axis == 1) { sum_evenly_spaced_rows_csr_fill_sparsity_and_idx_map( - x->jacobian, node->jacobian, node->d1, node->iwork, snode->idx_map); + x->jacobian, node->jacobian, node->size, node->iwork, snode->idx_map); } } @@ -157,33 +157,34 @@ static void free_type_data(expr *node) expr *new_sum(expr *child, int axis) { - int d1 = 0; + int d2 = 0; switch (axis) { case -1: /* no axis specified */ - d1 = 1; + d2 = 1; break; case 0: /* sum rows together */ - d1 = child->d2; + d2 = child->d2; break; case 1: /* sum columns together */ - d1 = child->d1; + d2 = child->d1; break; } /* Allocate the type-specific struct */ sum_expr *snode = (sum_expr *) calloc(1, sizeof(sum_expr)); expr *node = &snode->base; - init_expr(node, d1, 1, child->n_vars, forward, jacobian_init, eval_jacobian, + + /* to be consistent with CVXPY and NumPy we treat the result from + sum with an axis argument as a row vector */ + init_expr(node, 1, d2, child->n_vars, forward, jacobian_init, eval_jacobian, is_affine, free_type_data); node->left = child; expr_retain(child); - - /* hessian function pointers */ node->wsum_hess_init = wsum_hess_init; node->eval_wsum_hess = eval_wsum_hess; diff --git a/src/other/quad_form.c b/src/other/quad_form.c index da7fbae..9cb64dd 100644 --- a/src/other/quad_form.c +++ b/src/other/quad_form.c @@ -179,7 +179,7 @@ static void free_type_data(expr *node) expr *new_quad_form(expr *left, CSR_Matrix *Q) { - assert(left->d2 == 1); + assert(left->d1 == 1 || left->d2 == 1); /* left must be a vector */ quad_form_expr *qnode = (quad_form_expr *) calloc(1, sizeof(quad_form_expr)); expr *node = &qnode->base; From 7837da8cb9bcc79d50d5996b104a2ec33de0880a Mon Sep 17 00:00:00 2001 From: Daniel Date: Fri, 16 Jan 2026 05:13:55 -0800 Subject: [PATCH 3/8] edits --- src/affine/neg.c | 10 +++-- src/bivariate/const_scalar_mult.c | 3 ++ src/bivariate/quad_over_lin.c | 67 +++++++++++++++-------------- src/elementwise_univariate/common.c | 2 + src/other/prod.c | 2 + src/other/quad_form.c | 18 ++++---- src/problem.c | 1 + 7 files changed, 57 insertions(+), 46 deletions(-) diff --git a/src/affine/neg.c b/src/affine/neg.c index 632372b..71be271 100644 --- a/src/affine/neg.c +++ b/src/affine/neg.c @@ -1,4 +1,5 @@ #include "affine.h" +#include #include static void forward(expr *node, const double *u) @@ -43,17 +44,18 @@ static void eval_jacobian(expr *node) static void wsum_hess_init(expr *node) { + expr *x = node->left; + /* initialize child's wsum_hess */ - node->left->wsum_hess_init(node->left); + x->wsum_hess_init(x); /* same sparsity pattern as child */ - CSR_Matrix *child_hess = node->left->wsum_hess; - node->wsum_hess = new_csr_matrix(child_hess->m, child_hess->n, child_hess->nnz); + CSR_Matrix *child_hess = x->wsum_hess; + node->wsum_hess = new_csr_matrix(node->n_vars, node->n_vars, child_hess->nnz); /* copy row pointers and column indices (sparsity pattern is constant) */ memcpy(node->wsum_hess->p, child_hess->p, (child_hess->m + 1) * sizeof(int)); memcpy(node->wsum_hess->i, child_hess->i, child_hess->nnz * sizeof(int)); - node->wsum_hess->nnz = child_hess->nnz; } static void eval_wsum_hess(expr *node, const double *w) diff --git a/src/bivariate/const_scalar_mult.c b/src/bivariate/const_scalar_mult.c index b46be7c..dc08ec6 100644 --- a/src/bivariate/const_scalar_mult.c +++ b/src/bivariate/const_scalar_mult.c @@ -89,5 +89,8 @@ expr *new_const_scalar_mult(double a, expr *child) mult_node->a = a; expr_retain(child); + // just for debugging, should be removed + strcpy(node->name, "const_scalar_mult"); + return node; } diff --git a/src/bivariate/quad_over_lin.c b/src/bivariate/quad_over_lin.c index 1e36e29..47f5714 100644 --- a/src/bivariate/quad_over_lin.c +++ b/src/bivariate/quad_over_lin.c @@ -38,23 +38,23 @@ static void jacobian_init(expr *node) /* if left node is a variable */ if (x->var_id != NOT_A_VARIABLE) { - node->jacobian = new_csr_matrix(1, node->n_vars, x->d1 + 1); + node->jacobian = new_csr_matrix(1, node->n_vars, x->size + 1); node->jacobian->p[0] = 0; - node->jacobian->p[1] = x->d1 + 1; + node->jacobian->p[1] = x->size + 1; /* if x has lower idx than y*/ if (x->var_id < y->var_id) { - for (int j = 0; j < x->d1; j++) + for (int j = 0; j < x->size; j++) { node->jacobian->i[j] = x->var_id + j; } - node->jacobian->i[x->d1] = y->var_id; + node->jacobian->i[x->size] = y->var_id; } else /* y has lower idx than x */ { node->jacobian->i[0] = y->var_id; - for (int j = 0; j < x->d1; j++) + for (int j = 0; j < x->size; j++) { node->jacobian->i[j + 1] = x->var_id + j; } @@ -63,7 +63,7 @@ static void jacobian_init(expr *node) else /* left node is not a variable (guaranteed to be a linear operator) */ { linear_op_expr *lin_x = (linear_op_expr *) x; - node->dwork = (double *) malloc(x->d1 * sizeof(double)); + node->dwork = (double *) malloc(x->size * sizeof(double)); /* compute required allocation and allocate jacobian */ bool *col_nz = (bool *) calloc( @@ -115,16 +115,16 @@ static void eval_jacobian(expr *node) /* if x has lower idx than y*/ if (x->var_id < y->var_id) { - for (int j = 0; j < x->d1; j++) + for (int j = 0; j < x->size; j++) { node->jacobian->x[j] = (2.0 * x->value[j]) / y->value[0]; } - node->jacobian->x[x->d1] = -node->value[0] / y->value[0]; + node->jacobian->x[x->size] = -node->value[0] / y->value[0]; } else /* y has lower idx than x */ { node->jacobian->x[0] = -node->value[0] / y->value[0]; - for (int j = 0; j < x->d1; j++) + for (int j = 0; j < x->size; j++) { node->jacobian->x[j + 1] = (2.0 * x->value[j]) / y->value[0]; } @@ -135,7 +135,7 @@ static void eval_jacobian(expr *node) CSC_Matrix *A_csc = ((linear_op_expr *) x)->A_csc; /* local jacobian */ - for (int j = 0; j < x->d1; j++) + for (int j = 0; j < x->size; j++) { node->dwork[j] = (2.0 * x->value[j]) / y->value[0]; } @@ -160,14 +160,15 @@ static void wsum_hess_init(expr *node) /* if left node is a variable */ if (x->var_id != NOT_A_VARIABLE) { - node->wsum_hess = new_csr_matrix(node->n_vars, node->n_vars, 3 * x->d1 + 1); + node->wsum_hess = + new_csr_matrix(node->n_vars, node->n_vars, 3 * x->size + 1); CSR_Matrix *H = node->wsum_hess; /* if x has lower idx than y*/ if (var_id_x < var_id_y) { /* x rows: each row has 2 entries (diagonal element + element for y) */ - for (int i = 0; i < x->d1; i++) + for (int i = 0; i < x->size; i++) { H->p[var_id_x + i] = 2 * i; H->i[2 * i] = var_id_x + i; @@ -175,45 +176,45 @@ static void wsum_hess_init(expr *node) } /* rows between x and y are empty, all point to same offset */ - int offset = 2 * x->d1; - for (int i = var_id_x + x->d1; i <= var_id_y; i++) + int offset = 2 * x->size; + for (int i = var_id_x + x->size; i <= var_id_y; i++) { H->p[i] = offset; } /* y row has d1 + 1 entries */ - H->p[var_id_y + 1] = offset + x->d1 + 1; - for (int i = 0; i < x->d1; i++) + H->p[var_id_y + 1] = offset + x->size + 1; + for (int i = 0; i < x->size; i++) { H->i[offset + i] = var_id_x + i; } - H->i[offset + x->d1] = var_id_y; + H->i[offset + x->size] = var_id_y; /* remaining rows are empty */ for (int i = var_id_y + 1; i <= node->n_vars; i++) { - H->p[i] = 3 * x->d1 + 1; + H->p[i] = 3 * x->size + 1; } } else /* y has lower idx than x */ { /* y row has d1 + 1 entries */ - H->p[var_id_y + 1] = x->d1 + 1; + H->p[var_id_y + 1] = x->size + 1; H->i[0] = var_id_y; - for (int i = 0; i < x->d1; i++) + for (int i = 0; i < x->size; i++) { H->i[i + 1] = var_id_x + i; } /* rows between y and x are empty, all point to same offset */ - int offset = x->d1 + 1; + int offset = x->size + 1; for (int i = var_id_y + 1; i <= var_id_x; i++) { H->p[i] = offset; } /* x rows: each row has 2 entries */ - for (int i = 0; i < x->d1; i++) + for (int i = 0; i < x->size; i++) { H->p[var_id_x + i] = offset + 2 * i; H->i[offset + 2 * i] = var_id_y; @@ -221,9 +222,9 @@ static void wsum_hess_init(expr *node) } /* remaining rows are empty */ - for (int i = var_id_x + x->d1; i <= node->n_vars; i++) + for (int i = var_id_x + x->size; i <= node->n_vars; i++) { - H->p[i] = 3 * x->d1 + 1; + H->p[i] = 3 * x->size + 1; } } } @@ -241,7 +242,7 @@ static void eval_wsum_hess(expr *node, const double *w) double *H = node->wsum_hess->x; int var_id_x = node->left->var_id; int var_id_y = node->right->var_id; - int x_d1 = node->left->d1; + int x_size = node->left->size; double a = (2.0 * w[0]) / y; double b = -(2.0 * w[0]) / (y * y); @@ -252,32 +253,32 @@ static void eval_wsum_hess(expr *node, const double *w) if (var_id_x < var_id_y) { /* x rows*/ - for (int i = 0; i < x_d1; i++) + for (int i = 0; i < x_size; i++) { H[2 * i] = a; H[2 * i + 1] = b * x[i]; } /* y row */ - int offset = 2 * x_d1; - for (int i = 0; i < x_d1; i++) + int offset = 2 * x_size; + for (int i = 0; i < x_size; i++) { H[offset + i] = b * x[i]; } - H[offset + x_d1] = -b * node->value[0]; + H[offset + x_size] = -b * node->value[0]; } else /* y has lower idx than x */ { /* y row */ H[0] = -b * node->value[0]; - for (int i = 0; i < x_d1; i++) + for (int i = 0; i < x_size; i++) { H[i + 1] = b * x[i]; } /* x rows*/ - int offset = x_d1 + 1; - for (int i = 0; i < x_d1; i++) + int offset = x_size + 1; + for (int i = 0; i < x_size; i++) { H[offset + 2 * i] = b * x[i]; H[offset + 2 * i + 1] = a; @@ -296,7 +297,7 @@ static void eval_wsum_hess(expr *node, const double *w) expr *new_quad_over_lin(expr *left, expr *right) { - assert((left->d2 == 1 && right->d2 == 1)); /* right must be scalar*/ + assert((right->d2 == 1 && right->d2 == 1)); /* right must be scalar*/ expr *node = new_expr(1, 1, left->n_vars); node->left = left; node->right = right; diff --git a/src/elementwise_univariate/common.c b/src/elementwise_univariate/common.c index 2af054e..fb58f69 100644 --- a/src/elementwise_univariate/common.c +++ b/src/elementwise_univariate/common.c @@ -1,7 +1,9 @@ #include "elementwise_univariate.h" #include "expr.h" #include "subexpr.h" +#include #include +#include void jacobian_init_elementwise(expr *node) { diff --git a/src/other/prod.c b/src/other/prod.c index 92fbe13..35cbeb6 100644 --- a/src/other/prod.c +++ b/src/other/prod.c @@ -179,6 +179,8 @@ static void free_type_data(expr *node) (void) node; } +/* when we implement axis-support, check convention of numpy and cvxpy. +I think they return row vectors.*/ expr *new_prod(expr *child) { /* Output is scalar: 1 x 1 */ diff --git a/src/other/quad_form.c b/src/other/quad_form.c index 9cb64dd..42fb989 100644 --- a/src/other/quad_form.c +++ b/src/other/quad_form.c @@ -19,7 +19,7 @@ static void forward(expr *node, const double *u) csr_matvec(Q, x->value, node->dwork, 0); node->value[0] = 0.0; - for (int i = 0; i < x->d1; i++) + for (int i = 0; i < x->size; i++) { node->value[0] += x->value[i] * node->dwork[i]; } @@ -28,14 +28,14 @@ static void forward(expr *node, const double *u) static void jacobian_init(expr *node) { assert(node->left->var_id != NOT_A_VARIABLE); - assert(node->left->d2 == 1); + expr *x = node->left; - node->dwork = (double *) malloc(x->d1 * sizeof(double)); - node->jacobian = new_csr_matrix(1, node->n_vars, x->d1); + node->dwork = (double *) malloc(x->size * sizeof(double)); + node->jacobian = new_csr_matrix(1, node->n_vars, x->size); node->jacobian->p[0] = 0; - node->jacobian->p[1] = x->d1; + node->jacobian->p[1] = x->size; - for (int j = 0; j < x->d1; j++) + for (int j = 0; j < x->size; j++) { node->jacobian->i[j] = x->var_id + j; } @@ -49,7 +49,7 @@ static void eval_jacobian(expr *node) // jacobian = 2 * Q * x csr_matvec(Q, x->value, node->jacobian->x, 0); - for (int j = 0; j < x->d1; j++) + for (int j = 0; j < x->size; j++) { node->jacobian->x[j] *= 2.0; } @@ -62,8 +62,8 @@ static void wsum_hess_init(expr *node) CSR_Matrix *H = new_csr_matrix(node->n_vars, node->n_vars, Q->nnz); /* set global row pointers */ - memcpy(H->p + x->var_id, Q->p, (x->d1 + 1) * sizeof(int)); - for (int i = x->var_id + x->d1 + 1; i <= node->n_vars; i++) + memcpy(H->p + x->var_id, Q->p, (x->size + 1) * sizeof(int)); + for (int i = x->var_id + x->size + 1; i <= node->n_vars; i++) { H->p[i] = Q->nnz; diff --git a/src/problem.c b/src/problem.c index 5ac4634..0372494 100644 --- a/src/problem.c +++ b/src/problem.c @@ -1,6 +1,7 @@ #include "problem.h" #include "utils/utils.h" #include +#include #include #include From f0021cc40d759b0b3bd1b5c3b7fc5fd22f61ef72 Mon Sep 17 00:00:00 2001 From: Daniel Date: Fri, 16 Jan 2026 05:59:29 -0800 Subject: [PATCH 4/8] fixing rel entropy test --- python/atoms/rel_entr_scalar_vector.h | 37 ++++ python/atoms/rel_entr_vector_scalar.h | 37 ++++ python/bindings.c | 6 + src/bivariate/const_vector_mult.c | 1 - src/bivariate/rel_entr.c | 4 - src/bivariate/rel_entr_scalar_vector.c | 208 +++++++++++++++++ src/bivariate/rel_entr_vector_scalar.c | 209 ++++++++++++++++++ tests/all_tests.c | 8 + .../test_rel_entr_scalar_vector.h | 36 +++ .../test_rel_entr_vector_scalar.h | 36 +++ tests/wsum_hess/test_rel_entr_scalar_vector.h | 33 +++ tests/wsum_hess/test_rel_entr_vector_scalar.h | 33 +++ 12 files changed, 643 insertions(+), 5 deletions(-) create mode 100644 python/atoms/rel_entr_scalar_vector.h create mode 100644 python/atoms/rel_entr_vector_scalar.h create mode 100644 src/bivariate/rel_entr_scalar_vector.c create mode 100644 src/bivariate/rel_entr_vector_scalar.c create mode 100644 tests/jacobian_tests/test_rel_entr_scalar_vector.h create mode 100644 tests/jacobian_tests/test_rel_entr_vector_scalar.h create mode 100644 tests/wsum_hess/test_rel_entr_scalar_vector.h create mode 100644 tests/wsum_hess/test_rel_entr_vector_scalar.h diff --git a/python/atoms/rel_entr_scalar_vector.h b/python/atoms/rel_entr_scalar_vector.h new file mode 100644 index 0000000..85bcfe7 --- /dev/null +++ b/python/atoms/rel_entr_scalar_vector.h @@ -0,0 +1,37 @@ +// SPDX-License-Identifier: Apache-2.0 + +#ifndef ATOM_REL_ENTR_SCALAR_VECTOR_H +#define ATOM_REL_ENTR_SCALAR_VECTOR_H + +#include "bivariate.h" +#include "common.h" + +/* rel_entr_scalar_vector: rel_entr(x, y) where x is scalar, y is vector */ +static PyObject *py_make_rel_entr_scalar_vector(PyObject *self, PyObject *args) +{ + (void) self; + PyObject *left_capsule, *right_capsule; + if (!PyArg_ParseTuple(args, "OO", &left_capsule, &right_capsule)) + { + return NULL; + } + expr *left = (expr *) PyCapsule_GetPointer(left_capsule, EXPR_CAPSULE_NAME); + expr *right = (expr *) PyCapsule_GetPointer(right_capsule, EXPR_CAPSULE_NAME); + if (!left || !right) + { + PyErr_SetString(PyExc_ValueError, "invalid child capsule"); + return NULL; + } + + expr *node = new_rel_entr_first_arg_scalar(left, right); + if (!node) + { + PyErr_SetString(PyExc_RuntimeError, + "failed to create rel_entr_scalar_vector node"); + return NULL; + } + expr_retain(node); /* Capsule owns a reference */ + return PyCapsule_New(node, EXPR_CAPSULE_NAME, expr_capsule_destructor); +} + +#endif /* ATOM_REL_ENTR_SCALAR_VECTOR_H */ diff --git a/python/atoms/rel_entr_vector_scalar.h b/python/atoms/rel_entr_vector_scalar.h new file mode 100644 index 0000000..d671646 --- /dev/null +++ b/python/atoms/rel_entr_vector_scalar.h @@ -0,0 +1,37 @@ +// SPDX-License-Identifier: Apache-2.0 + +#ifndef ATOM_REL_ENTR_VECTOR_SCALAR_H +#define ATOM_REL_ENTR_VECTOR_SCALAR_H + +#include "bivariate.h" +#include "common.h" + +/* rel_entr_vector_scalar: rel_entr(x, y) where x is vector, y is scalar */ +static PyObject *py_make_rel_entr_vector_scalar(PyObject *self, PyObject *args) +{ + (void) self; + PyObject *left_capsule, *right_capsule; + if (!PyArg_ParseTuple(args, "OO", &left_capsule, &right_capsule)) + { + return NULL; + } + expr *left = (expr *) PyCapsule_GetPointer(left_capsule, EXPR_CAPSULE_NAME); + expr *right = (expr *) PyCapsule_GetPointer(right_capsule, EXPR_CAPSULE_NAME); + if (!left || !right) + { + PyErr_SetString(PyExc_ValueError, "invalid child capsule"); + return NULL; + } + + expr *node = new_rel_entr_second_arg_scalar(left, right); + if (!node) + { + PyErr_SetString(PyExc_RuntimeError, + "failed to create rel_entr_vector_scalar node"); + return NULL; + } + expr_retain(node); /* Capsule owns a reference */ + return PyCapsule_New(node, EXPR_CAPSULE_NAME, expr_capsule_destructor); +} + +#endif /* ATOM_REL_ENTR_VECTOR_SCALAR_H */ diff --git a/python/bindings.c b/python/bindings.c index a6d946e..12a2eea 100644 --- a/python/bindings.c +++ b/python/bindings.c @@ -27,6 +27,8 @@ #include "atoms/quad_form.h" #include "atoms/quad_over_lin.h" #include "atoms/rel_entr.h" +#include "atoms/rel_entr_scalar_vector.h" +#include "atoms/rel_entr_vector_scalar.h" #include "atoms/reshape.h" #include "atoms/right_matmul.h" #include "atoms/sin.h" @@ -96,6 +98,10 @@ static PyMethodDef DNLPMethods[] = { "Create quad_over_lin node (sum(x^2) / y)"}, {"make_rel_entr", py_make_rel_entr, METH_VARARGS, "Create rel_entr node: x * log(x/y) elementwise"}, + {"make_rel_entr_vector_scalar", py_make_rel_entr_vector_scalar, METH_VARARGS, + "Create rel_entr node with vector first arg, scalar second arg"}, + {"make_rel_entr_scalar_vector", py_make_rel_entr_scalar_vector, METH_VARARGS, + "Create rel_entr node with scalar first arg, vector second arg"}, {"get_expr_dimensions", py_get_expr_dimensions, METH_VARARGS, "Get the dimensions (d1, d2) of an expression"}, {"get_expr_size", py_get_expr_size, METH_VARARGS, diff --git a/src/bivariate/const_vector_mult.c b/src/bivariate/const_vector_mult.c index 0f3f921..4846a7c 100644 --- a/src/bivariate/const_vector_mult.c +++ b/src/bivariate/const_vector_mult.c @@ -23,7 +23,6 @@ static void forward(expr *node, const double *u) static void jacobian_init(expr *node) { - printf("jacobian_init const_vector_mult\n \n \n"); expr *x = node->left; /* initialize child jacobian */ diff --git a/src/bivariate/rel_entr.c b/src/bivariate/rel_entr.c index 61be542..f0496df 100644 --- a/src/bivariate/rel_entr.c +++ b/src/bivariate/rel_entr.c @@ -183,7 +183,3 @@ expr *new_rel_entr_vector_args(expr *left, expr *right) node->eval_wsum_hess = eval_wsum_hess_vector_args; return node; } - -// -------------------------------------------------------------------- -// Implementation of relative entropy when one argument is a vector -// -------------------------------------------------------------------- diff --git a/src/bivariate/rel_entr_scalar_vector.c b/src/bivariate/rel_entr_scalar_vector.c new file mode 100644 index 0000000..1d0758e --- /dev/null +++ b/src/bivariate/rel_entr_scalar_vector.c @@ -0,0 +1,208 @@ +#include "bivariate.h" +#include +#include +#include +#include + +// -------------------------------------------------------------------- +// Implementation of relative entropy when first argument is a scalar +// and second argument is a vector. +// -------------------------------------------------------------------- +static void forward_scalar_vector(expr *node, const double *u) +{ + expr *x = node->left; + expr *y = node->right; + + /* children's forward passes */ + x->forward(x, u); + y->forward(y, u); + + /* local forward pass */ + for (int i = 0; i < node->size; i++) + { + node->value[i] = x->value[0] * log(x->value[0] / y->value[i]); + } +} + +static void jacobian_init_scalar_vector(expr *node) +{ + expr *x = node->left; + expr *y = node->right; + assert(x->d1 == 1 && x->d2 == 1); + assert(x->var_id != NOT_A_VARIABLE && y->var_id != NOT_A_VARIABLE); + assert(x->var_id != y->var_id); + + node->jacobian = new_csr_matrix(node->size, node->n_vars, 2 * node->size); + + if (x->var_id < y->var_id) + { + for (int j = 0; j < node->size; j++) + { + node->jacobian->i[2 * j] = x->var_id; + node->jacobian->i[2 * j + 1] = y->var_id + j; + node->jacobian->p[j] = 2 * j; + } + } + else + { + for (int j = 0; j < node->size; j++) + { + node->jacobian->i[2 * j] = y->var_id + j; + node->jacobian->i[2 * j + 1] = x->var_id; + node->jacobian->p[j] = 2 * j; + } + } + + node->jacobian->p[node->size] = 2 * node->size; +} + +static void eval_jacobian_scalar_vector(expr *node) +{ + expr *x = node->left; + expr *y = node->right; + + if (x->var_id < y->var_id) + { + for (int i = 0; i < node->size; i++) + { + node->jacobian->x[2 * i] = log(x->value[0] / y->value[i]) + 1; + node->jacobian->x[2 * i + 1] = -x->value[0] / y->value[i]; + } + } + else + { + for (int i = 0; i < node->size; i++) + { + node->jacobian->x[2 * i] = -x->value[0] / y->value[i]; + node->jacobian->x[2 * i + 1] = log(x->value[0] / y->value[i]) + 1; + } + } +} + +static void wsum_hess_init_scalar_vector(expr *node) +{ + expr *x = node->left; + expr *y = node->right; + int var_id_x = x->var_id; + int var_id_y = y->var_id; + + node->wsum_hess = new_csr_matrix(node->n_vars, node->n_vars, 3 * node->size + 1); + CSR_Matrix *H = node->wsum_hess; + + if (var_id_x < var_id_y) + { + H->p[var_id_x + 1] = node->size + 1; + H->i[0] = var_id_x; + for (int i = 0; i < node->size; i++) + { + H->i[i + 1] = var_id_y + i; + } + + int offset = node->size + 1; + for (int i = var_id_x + 1; i <= var_id_y; i++) + { + H->p[i] = offset; + } + + for (int i = 0; i < node->size; i++) + { + H->p[var_id_y + i] = offset + 2 * i; + H->i[offset + 2 * i] = var_id_x; + H->i[offset + 2 * i + 1] = var_id_y + i; + } + + for (int i = var_id_y + node->size; i <= node->n_vars; i++) + { + H->p[i] = 3 * node->size + 1; + } + } + else + { + for (int i = 0; i < node->size; i++) + { + H->p[var_id_y + i] = 2 * i; + H->i[2 * i] = var_id_y + i; + H->i[2 * i + 1] = var_id_x; + } + + int offset = 2 * node->size; + for (int i = var_id_y + node->size; i <= var_id_x; i++) + { + H->p[i] = offset; + } + + H->p[var_id_x + 1] = offset + node->size + 1; + for (int i = 0; i < node->size; i++) + { + H->i[offset + i] = var_id_y + i; + } + H->i[offset + node->size] = var_id_x; + + for (int i = var_id_x + 1; i <= node->n_vars; i++) + { + H->p[i] = 3 * node->size + 1; + } + } +} + +static void eval_wsum_hess_scalar_vector(expr *node, const double *w) +{ + double x = node->left->value[0]; + double *y = node->right->value; + double *H = node->wsum_hess->x; + int var_id_x = node->left->var_id; + int var_id_y = node->right->var_id; + + double diag_x = 0.0; + for (int i = 0; i < node->size; i++) + { + diag_x += w[i] / x; + } + + if (var_id_x < var_id_y) + { + H[0] = diag_x; + for (int i = 0; i < node->size; i++) + { + H[i + 1] = -w[i] / y[i]; + } + + int offset = node->size + 1; + for (int i = 0; i < node->size; i++) + { + H[offset + 2 * i] = -w[i] / y[i]; + H[offset + 2 * i + 1] = w[i] * x / (y[i] * y[i]); + } + } + else + { + for (int i = 0; i < node->size; i++) + { + H[2 * i] = w[i] * x / (y[i] * y[i]); + H[2 * i + 1] = -w[i] / y[i]; + } + + int offset = 2 * node->size; + for (int i = 0; i < node->size; i++) + { + H[offset + i] = -w[i] / y[i]; + } + H[offset + node->size] = diag_x; + } +} + +expr *new_rel_entr_first_arg_scalar(expr *left, expr *right) +{ + assert(left->d1 == 1 && left->d2 == 1); + expr *node = new_expr(right->d1, right->d2, left->n_vars); + node->left = left; + node->right = right; + expr_retain(left); + expr_retain(right); + node->forward = forward_scalar_vector; + node->jacobian_init = jacobian_init_scalar_vector; + node->eval_jacobian = eval_jacobian_scalar_vector; + node->wsum_hess_init = wsum_hess_init_scalar_vector; + node->eval_wsum_hess = eval_wsum_hess_scalar_vector; + return node; +} diff --git a/src/bivariate/rel_entr_vector_scalar.c b/src/bivariate/rel_entr_vector_scalar.c new file mode 100644 index 0000000..0b36ca6 --- /dev/null +++ b/src/bivariate/rel_entr_vector_scalar.c @@ -0,0 +1,209 @@ +#include "bivariate.h" +#include +#include +#include +#include + +// -------------------------------------------------------------------- +// Implementation of relative entropy when first argument is a vector +// and second argument is a scalar. +// -------------------------------------------------------------------- +static void forward_vector_scalar(expr *node, const double *u) +{ + expr *x = node->left; + expr *y = node->right; + + /* children's forward passes */ + x->forward(x, u); + y->forward(y, u); + + /* local forward pass */ + for (int i = 0; i < node->size; i++) + { + node->value[i] = x->value[i] * log(x->value[i] / y->value[0]); + } +} + +static void jacobian_init_vector_scalar(expr *node) +{ + expr *x = node->left; + expr *y = node->right; + assert(y->d1 == 1 && y->d2 == 1); + assert(x->var_id != NOT_A_VARIABLE && y->var_id != NOT_A_VARIABLE); + assert(x->var_id != y->var_id); + + node->jacobian = new_csr_matrix(node->size, node->n_vars, 2 * node->size); + + if (x->var_id < y->var_id) + { + for (int j = 0; j < node->size; j++) + { + node->jacobian->i[2 * j] = x->var_id + j; + node->jacobian->i[2 * j + 1] = y->var_id; + node->jacobian->p[j] = 2 * j; + } + } + else + { + for (int j = 0; j < node->size; j++) + { + node->jacobian->i[2 * j] = y->var_id; + node->jacobian->i[2 * j + 1] = x->var_id + j; + node->jacobian->p[j] = 2 * j; + } + } + + node->jacobian->p[node->size] = 2 * node->size; +} + +static void eval_jacobian_vector_scalar(expr *node) +{ + expr *x = node->left; + expr *y = node->right; + + if (x->var_id < y->var_id) + { + for (int i = 0; i < node->size; i++) + { + node->jacobian->x[2 * i] = log(x->value[i] / y->value[0]) + 1; + node->jacobian->x[2 * i + 1] = -x->value[i] / y->value[0]; + } + } + else + { + for (int i = 0; i < node->size; i++) + { + node->jacobian->x[2 * i] = -x->value[i] / y->value[0]; + node->jacobian->x[2 * i + 1] = log(x->value[i] / y->value[0]) + 1; + } + } +} + +static void wsum_hess_init_vector_scalar(expr *node) +{ + expr *x = node->left; + expr *y = node->right; + int var_id_x = x->var_id; + int var_id_y = y->var_id; + + node->wsum_hess = new_csr_matrix(node->n_vars, node->n_vars, 3 * node->size + 1); + CSR_Matrix *H = node->wsum_hess; + + if (var_id_x < var_id_y) + { + for (int i = 0; i < node->size; i++) + { + H->p[var_id_x + i] = 2 * i; + H->i[2 * i] = var_id_x + i; + H->i[2 * i + 1] = var_id_y; + } + + int offset = 2 * node->size; + for (int i = var_id_x + node->size; i <= var_id_y; i++) + { + H->p[i] = offset; + } + + H->p[var_id_y + 1] = offset + node->size + 1; + for (int i = 0; i < node->size; i++) + { + H->i[offset + i] = var_id_x + i; + } + H->i[offset + node->size] = var_id_y; + + for (int i = var_id_y + 1; i <= node->n_vars; i++) + { + H->p[i] = 3 * node->size + 1; + } + } + else + { + H->p[var_id_y + 1] = node->size + 1; + H->i[0] = var_id_y; + for (int i = 0; i < node->size; i++) + { + H->i[i + 1] = var_id_x + i; + } + + int offset = node->size + 1; + for (int i = var_id_y + 1; i <= var_id_x; i++) + { + H->p[i] = offset; + } + + for (int i = 0; i < node->size; i++) + { + H->p[var_id_x + i] = offset + 2 * i; + H->i[offset + 2 * i] = var_id_y; + H->i[offset + 2 * i + 1] = var_id_x + i; + } + + for (int i = var_id_x + node->size; i <= node->n_vars; i++) + { + H->p[i] = 3 * node->size + 1; + } + } +} + +static void eval_wsum_hess_vector_scalar(expr *node, const double *w) +{ + double *x = node->left->value; + double y = node->right->value[0]; + double *H = node->wsum_hess->x; + int var_id_x = node->left->var_id; + int var_id_y = node->right->var_id; + + double diag_y = 0.0; + for (int i = 0; i < node->size; i++) + { + diag_y += w[i] * x[i]; + } + diag_y /= (y * y); + + if (var_id_x < var_id_y) + { + for (int i = 0; i < node->size; i++) + { + H[2 * i] = w[i] / x[i]; + H[2 * i + 1] = -w[i] / y; + } + + int offset = 2 * node->size; + for (int i = 0; i < node->size; i++) + { + H[offset + i] = -w[i] / y; + } + H[offset + node->size] = diag_y; + } + else + { + H[0] = diag_y; + for (int i = 0; i < node->size; i++) + { + H[i + 1] = -w[i] / y; + } + + int offset = node->size + 1; + for (int i = 0; i < node->size; i++) + { + H[offset + 2 * i] = -w[i] / y; + H[offset + 2 * i + 1] = w[i] / x[i]; + } + } +} + +expr *new_rel_entr_second_arg_scalar(expr *left, expr *right) +{ + assert(right->d1 == 1 && right->d2 == 1); + expr *node = new_expr(left->d1, left->d2, left->n_vars); + node->left = left; + node->right = right; + expr_retain(left); + expr_retain(right); + node->forward = forward_vector_scalar; + node->jacobian_init = jacobian_init_vector_scalar; + node->eval_jacobian = eval_jacobian_vector_scalar; + node->wsum_hess_init = wsum_hess_init_vector_scalar; + node->eval_wsum_hess = eval_wsum_hess_vector_scalar; + return node; +} diff --git a/tests/all_tests.c b/tests/all_tests.c index 0d0722d..0f8b9f6 100644 --- a/tests/all_tests.c +++ b/tests/all_tests.c @@ -29,6 +29,8 @@ #include "jacobian_tests/test_quad_form.h" #include "jacobian_tests/test_quad_over_lin.h" #include "jacobian_tests/test_rel_entr.h" +#include "jacobian_tests/test_rel_entr_scalar_vector.h" +#include "jacobian_tests/test_rel_entr_vector_scalar.h" #include "jacobian_tests/test_right_matmul.h" #include "jacobian_tests/test_sum.h" #include "jacobian_tests/test_trace.h" @@ -54,6 +56,8 @@ #include "wsum_hess/test_quad_form.h" #include "wsum_hess/test_quad_over_lin.h" #include "wsum_hess/test_rel_entr.h" +#include "wsum_hess/test_rel_entr_scalar_vector.h" +#include "wsum_hess/test_rel_entr_vector_scalar.h" #include "wsum_hess/test_right_matmul.h" #include "wsum_hess/test_sum.h" #include "wsum_hess/test_trace.h" @@ -97,6 +101,8 @@ int main(void) mu_run_test(test_jacobian_rel_entr_vector_args_1, tests_run); mu_run_test(test_jacobian_rel_entr_vector_args_2, tests_run); mu_run_test(test_jacobian_rel_entr_matrix_args, tests_run); + mu_run_test(test_jacobian_rel_entr_vector_scalar, tests_run); + mu_run_test(test_jacobian_rel_entr_scalar_vector, tests_run); mu_run_test(test_jacobian_elementwise_mult_1, tests_run); mu_run_test(test_jacobian_elementwise_mult_2, tests_run); mu_run_test(test_jacobian_elementwise_mult_3, tests_run); @@ -165,6 +171,8 @@ int main(void) mu_run_test(test_wsum_hess_rel_entr_1, tests_run); mu_run_test(test_wsum_hess_rel_entr_2, tests_run); mu_run_test(test_wsum_hess_rel_entr_matrix, tests_run); + mu_run_test(test_wsum_hess_rel_entr_vector_scalar, tests_run); + 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_index_log, tests_run); diff --git a/tests/jacobian_tests/test_rel_entr_scalar_vector.h b/tests/jacobian_tests/test_rel_entr_scalar_vector.h new file mode 100644 index 0000000..2f5b402 --- /dev/null +++ b/tests/jacobian_tests/test_rel_entr_scalar_vector.h @@ -0,0 +1,36 @@ +#include "bivariate.h" +#include "expr.h" +#include "minunit.h" +#include "test_helpers.h" +#include +#include + +const char *test_jacobian_rel_entr_scalar_vector() +{ + // x is scalar 1, y is 3x1 with values [1, 2, 3] + double u_vals[4] = {1.0, 1.0, 2.0, 3.0}; + expr *x = new_variable(1, 1, 0, 4); + expr *y = new_variable(3, 1, 1, 4); + expr *node = new_rel_entr_first_arg_scalar(x, y); + + node->forward(node, u_vals); + node->jacobian_init(node); + node->eval_jacobian(node); + + double a = log(1.0 / 1.0) + 1.0; + double b = log(1.0 / 2.0) + 1.0; + double c = log(1.0 / 3.0) + 1.0; + double d = -1.0 / 1.0; + double e = -1.0 / 2.0; + double f = -1.0 / 3.0; + + double expected_Ax[6] = {a, d, b, e, c, f}; + int expected_Ap[4] = {0, 2, 4, 6}; + int expected_Ai[6] = {0, 1, 0, 2, 0, 3}; + + mu_assert("vals fail", cmp_double_array(node->jacobian->x, expected_Ax, 6)); + mu_assert("rows fail", cmp_int_array(node->jacobian->p, expected_Ap, 4)); + mu_assert("cols fail", cmp_int_array(node->jacobian->i, expected_Ai, 6)); + free_expr(node); + return 0; +} diff --git a/tests/jacobian_tests/test_rel_entr_vector_scalar.h b/tests/jacobian_tests/test_rel_entr_vector_scalar.h new file mode 100644 index 0000000..6ed7a42 --- /dev/null +++ b/tests/jacobian_tests/test_rel_entr_vector_scalar.h @@ -0,0 +1,36 @@ +#include "bivariate.h" +#include "expr.h" +#include "minunit.h" +#include "test_helpers.h" +#include +#include + +const char *test_jacobian_rel_entr_vector_scalar() +{ + // x is 3x1 with values [1, 2, 3], y is scalar 4 + double u_vals[4] = {1.0, 2.0, 3.0, 4.0}; + expr *x = new_variable(3, 1, 0, 4); + expr *y = new_variable(1, 1, 3, 4); + expr *node = new_rel_entr_second_arg_scalar(x, y); + + node->forward(node, u_vals); + node->jacobian_init(node); + node->eval_jacobian(node); + + double a = log(1.0 / 4.0) + 1.0; + double b = log(2.0 / 4.0) + 1.0; + double c = log(3.0 / 4.0) + 1.0; + double d = -1.0 / 4.0; + double e = -2.0 / 4.0; + double f = -3.0 / 4.0; + + double expected_Ax[6] = {a, d, b, e, c, f}; + int expected_Ap[4] = {0, 2, 4, 6}; + int expected_Ai[6] = {0, 3, 1, 3, 2, 3}; + + mu_assert("vals fail", cmp_double_array(node->jacobian->x, expected_Ax, 6)); + mu_assert("rows fail", cmp_int_array(node->jacobian->p, expected_Ap, 4)); + mu_assert("cols fail", cmp_int_array(node->jacobian->i, expected_Ai, 6)); + free_expr(node); + return 0; +} diff --git a/tests/wsum_hess/test_rel_entr_scalar_vector.h b/tests/wsum_hess/test_rel_entr_scalar_vector.h new file mode 100644 index 0000000..3e1db8a --- /dev/null +++ b/tests/wsum_hess/test_rel_entr_scalar_vector.h @@ -0,0 +1,33 @@ +#include "bivariate.h" +#include "expr.h" +#include "minunit.h" +#include "test_helpers.h" +#include +#include + +const char *test_wsum_hess_rel_entr_scalar_vector() +{ + // x is scalar 1, y is 3x1 with values [2, 3, 4], w = [4, 5, 6] + double u_vals[4] = {1.0, 2.0, 3.0, 4.0}; + double w[3] = {4.0, 5.0, 6.0}; + + expr *x = new_variable(1, 1, 0, 4); + expr *y = new_variable(3, 1, 1, 4); + expr *node = new_rel_entr_first_arg_scalar(x, y); + + node->forward(node, u_vals); + node->wsum_hess_init(node); + node->eval_wsum_hess(node, w); + + int expected_p[5] = {0, 4, 6, 8, 10}; + int expected_i[10] = {0, 1, 2, 3, 0, 1, 0, 2, 0, 3}; + double expected_x[10] = {15.0, -2.0, -5.0 / 3.0, -1.5, -2.0, + 1.0, -5.0 / 3.0, 5.0 / 9.0, -1.5, 0.375}; + + mu_assert("p array fails", cmp_int_array(node->wsum_hess->p, expected_p, 5)); + mu_assert("i array fails", cmp_int_array(node->wsum_hess->i, expected_i, 10)); + mu_assert("x array fails", cmp_double_array(node->wsum_hess->x, expected_x, 10)); + + free_expr(node); + return 0; +} diff --git a/tests/wsum_hess/test_rel_entr_vector_scalar.h b/tests/wsum_hess/test_rel_entr_vector_scalar.h new file mode 100644 index 0000000..83e6a53 --- /dev/null +++ b/tests/wsum_hess/test_rel_entr_vector_scalar.h @@ -0,0 +1,33 @@ +#include "bivariate.h" +#include "expr.h" +#include "minunit.h" +#include "test_helpers.h" +#include +#include + +const char *test_wsum_hess_rel_entr_vector_scalar() +{ + // x is 3x1 with values [1, 2, 3], y is scalar 4, w = [1, 2, 3] + double u_vals[4] = {1.0, 2.0, 3.0, 4.0}; + double w[3] = {1.0, 2.0, 3.0}; + + expr *x = new_variable(3, 1, 0, 4); + expr *y = new_variable(1, 1, 3, 4); + expr *node = new_rel_entr_second_arg_scalar(x, y); + + node->forward(node, u_vals); + node->wsum_hess_init(node); + node->eval_wsum_hess(node, w); + + int expected_p[5] = {0, 2, 4, 6, 10}; + int expected_i[10] = {0, 3, 1, 3, 2, 3, 0, 1, 2, 3}; + double expected_x[10] = {1.0, -0.25, 1.0, -0.5, 1.0, + -0.75, -0.25, -0.5, -0.75, 0.875}; + + mu_assert("p array fails", cmp_int_array(node->wsum_hess->p, expected_p, 5)); + mu_assert("i array fails", cmp_int_array(node->wsum_hess->i, expected_i, 10)); + mu_assert("x array fails", cmp_double_array(node->wsum_hess->x, expected_x, 10)); + + free_expr(node); + return 0; +} From 6d6adb9d7e5f1b03a59c2db0f97b61b18e3da697 Mon Sep 17 00:00:00 2001 From: Daniel Date: Fri, 16 Jan 2026 09:12:57 -0800 Subject: [PATCH 5/8] fixed damm index buggit add ../src/affine/index.c --- src/affine/index.c | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/src/affine/index.c b/src/affine/index.c index c8f6ade..49948a9 100644 --- a/src/affine/index.c +++ b/src/affine/index.c @@ -1,5 +1,7 @@ #include "affine.h" #include "subexpr.h" +#include +#include #include #include @@ -57,6 +59,7 @@ static void jacobian_init(expr *node) J->p[i + 1] = J->p[i] + len; } + J->nnz = J->p[idx->n_idxs]; node->jacobian = J; } @@ -142,16 +145,14 @@ static void free_type_data(expr *node) } } -expr *new_index(expr *child, const int *indices, int n_idxs) +expr *new_index(expr *child, int d1, int d2, const int *indices, int n_idxs) { + assert(d1 * d2 == n_idxs); /* allocate type-specific struct */ index_expr *idx = (index_expr *) calloc(1, sizeof(index_expr)); expr *node = &idx->base; - /* to be consistent with CVXPY and NumPy we treat the result from - indexing as a row vector. - TODO: is this correct? What if the index in numpy/cvxpy returns a matrix? */ - init_expr(node, 1, n_idxs, child->n_vars, forward, jacobian_init, eval_jacobian, + init_expr(node, d1, d2, child->n_vars, forward, jacobian_init, eval_jacobian, is_affine, free_type_data); node->left = child; From c6d974a2c8d63b3840d6913fceb356740922d065 Mon Sep 17 00:00:00 2001 From: Daniel Date: Fri, 16 Jan 2026 13:07:12 -0800 Subject: [PATCH 6/8] index debug --- include/affine.h | 2 +- include/subexpr.h | 2 - python/atoms/index.h | 5 ++- python/bindings.c | 4 ++ python/problem/hessian.h | 49 +++++++++++++++++++- python/problem/jacobian.h | 46 +++++++++++++++++++ src/affine/add.c | 1 + src/affine/broadcast.c | 74 +++++++++++++++---------------- src/affine/index.c | 3 +- src/affine/neg.c | 11 +++-- src/affine/reshape.c | 19 ++++---- tests/jacobian_tests/test_index.h | 12 ++--- tests/wsum_hess/test_index.h | 24 ++++++---- 13 files changed, 174 insertions(+), 78 deletions(-) diff --git a/include/affine.h b/include/affine.h index ec99717..a4c65d7 100644 --- a/include/affine.h +++ b/include/affine.h @@ -18,7 +18,7 @@ expr *new_trace(expr *child); expr *new_constant(int d1, int d2, int n_vars, const double *values); expr *new_variable(int d1, int d2, int var_id, int n_vars); -expr *new_index(expr *child, const int *indices, int n_idxs); +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); diff --git a/include/subexpr.h b/include/subexpr.h index 66df2ce..fd0a5b2 100644 --- a/include/subexpr.h +++ b/include/subexpr.h @@ -130,8 +130,6 @@ typedef struct broadcast_expr { expr base; broadcast_type type; - int m; /* target rows */ - int n; /* target cols */ } broadcast_expr; #endif /* SUBEXPR_H */ diff --git a/python/atoms/index.h b/python/atoms/index.h index 035ac56..64125ad 100644 --- a/python/atoms/index.h +++ b/python/atoms/index.h @@ -11,9 +11,10 @@ static PyObject *py_make_index(PyObject *self, PyObject *args) { PyObject *child_capsule; + int d1, d2; PyObject *indices_obj; - if (!PyArg_ParseTuple(args, "OO", &child_capsule, &indices_obj)) + if (!PyArg_ParseTuple(args, "OiiO", &child_capsule, &d1, &d2, &indices_obj)) { return NULL; } @@ -37,7 +38,7 @@ static PyObject *py_make_index(PyObject *self, PyObject *args) int n_idxs = (int) PyArray_SIZE(indices_array); int *indices_data = (int *) PyArray_DATA(indices_array); - expr *node = new_index(child, indices_data, n_idxs); + expr *node = new_index(child, d1, d2, indices_data, n_idxs); Py_DECREF(indices_array); diff --git a/python/bindings.c b/python/bindings.c index 12a2eea..55fbdc2 100644 --- a/python/bindings.c +++ b/python/bindings.c @@ -119,8 +119,12 @@ static PyMethodDef DNLPMethods[] = { "Compute objective gradient"}, {"problem_jacobian", py_problem_jacobian, METH_VARARGS, "Compute constraint jacobian"}, + {"get_jacobian", py_get_jacobian, METH_VARARGS, + "Get constraint jacobian without recomputing"}, {"problem_hessian", py_problem_hessian, METH_VARARGS, "Compute Lagrangian Hessian"}, + {"get_hessian", py_get_hessian, METH_VARARGS, + "Get Lagrangian Hessian without recomputing"}, {NULL, NULL, 0, NULL}}; static struct PyModuleDef dnlp_module = { diff --git a/python/problem/hessian.h b/python/problem/hessian.h index 41c2788..ea7f226 100644 --- a/python/problem/hessian.h +++ b/python/problem/hessian.h @@ -9,7 +9,8 @@ * Args: * prob_capsule: PyCapsule containing problem pointer * obj_factor: Scaling factor for objective Hessian (double) - * lagrange: Array of Lagrange multipliers (numpy array, length = total_constraint_size) + * lagrange: Array of Lagrange multipliers (numpy array, length = + * total_constraint_size) * * Returns: * Tuple of (data, indices, indptr, (m, n)) for scipy.sparse.csr_matrix @@ -73,4 +74,50 @@ static PyObject *py_problem_hessian(PyObject *self, PyObject *args) return Py_BuildValue("(OOO(ii))", data, indices, indptr, H->m, H->n); } +static PyObject *py_get_hessian(PyObject *self, PyObject *args) +{ + PyObject *prob_capsule; + if (!PyArg_ParseTuple(args, "O", &prob_capsule)) + { + return NULL; + } + + problem *prob = + (problem *) PyCapsule_GetPointer(prob_capsule, PROBLEM_CAPSULE_NAME); + if (!prob) + { + PyErr_SetString(PyExc_ValueError, "invalid problem capsule"); + return NULL; + } + + if (!prob->lagrange_hessian) + { + PyErr_SetString(PyExc_RuntimeError, + "hessian not initialized - call problem_hessian first"); + return NULL; + } + + CSR_Matrix *H = prob->lagrange_hessian; + npy_intp nnz = H->nnz; + npy_intp n_plus_1 = H->n + 1; + + PyObject *data = PyArray_SimpleNew(1, &nnz, NPY_DOUBLE); + PyObject *indices = PyArray_SimpleNew(1, &nnz, NPY_INT32); + PyObject *indptr = PyArray_SimpleNew(1, &n_plus_1, NPY_INT32); + + if (!data || !indices || !indptr) + { + Py_XDECREF(data); + Py_XDECREF(indices); + Py_XDECREF(indptr); + return NULL; + } + + memcpy(PyArray_DATA((PyArrayObject *) data), H->x, nnz * sizeof(double)); + memcpy(PyArray_DATA((PyArrayObject *) indices), H->i, nnz * sizeof(int)); + memcpy(PyArray_DATA((PyArrayObject *) indptr), H->p, n_plus_1 * sizeof(int)); + + return Py_BuildValue("(OOO(ii))", data, indices, indptr, H->m, H->n); +} + #endif /* PROBLEM_HESSIAN_H */ diff --git a/python/problem/jacobian.h b/python/problem/jacobian.h index 77fe795..1ba0a42 100644 --- a/python/problem/jacobian.h +++ b/python/problem/jacobian.h @@ -56,4 +56,50 @@ static PyObject *py_problem_jacobian(PyObject *self, PyObject *args) return Py_BuildValue("(OOO(ii))", data, indices, indptr, jac->m, jac->n); } +static PyObject *py_get_jacobian(PyObject *self, PyObject *args) +{ + PyObject *prob_capsule; + if (!PyArg_ParseTuple(args, "O", &prob_capsule)) + { + return NULL; + } + + problem *prob = + (problem *) PyCapsule_GetPointer(prob_capsule, PROBLEM_CAPSULE_NAME); + if (!prob) + { + PyErr_SetString(PyExc_ValueError, "invalid problem capsule"); + return NULL; + } + + if (!prob->jacobian) + { + PyErr_SetString(PyExc_RuntimeError, + "jacobian not initialized - call problem_jacobian first"); + return NULL; + } + + CSR_Matrix *jac = prob->jacobian; + npy_intp nnz = jac->nnz; + npy_intp m_plus_1 = jac->m + 1; + + PyObject *data = PyArray_SimpleNew(1, &nnz, NPY_DOUBLE); + PyObject *indices = PyArray_SimpleNew(1, &nnz, NPY_INT32); + PyObject *indptr = PyArray_SimpleNew(1, &m_plus_1, NPY_INT32); + + if (!data || !indices || !indptr) + { + Py_XDECREF(data); + Py_XDECREF(indices); + Py_XDECREF(indptr); + return NULL; + } + + memcpy(PyArray_DATA((PyArrayObject *) data), jac->x, nnz * sizeof(double)); + memcpy(PyArray_DATA((PyArrayObject *) indices), jac->i, nnz * sizeof(int)); + memcpy(PyArray_DATA((PyArrayObject *) indptr), jac->p, m_plus_1 * sizeof(int)); + + return Py_BuildValue("(OOO(ii))", data, indices, indptr, jac->m, jac->n); +} + #endif /* PROBLEM_JACOBIAN_H */ diff --git a/src/affine/add.c b/src/affine/add.c index 9f6549b..22ece09 100644 --- a/src/affine/add.c +++ b/src/affine/add.c @@ -1,5 +1,6 @@ #include "affine.h" #include +#include static void forward(expr *node, const double *u) { diff --git a/src/affine/broadcast.c b/src/affine/broadcast.c index a26bd35..d614db9 100644 --- a/src/affine/broadcast.c +++ b/src/affine/broadcast.c @@ -23,20 +23,20 @@ static void forward(expr *node, const double *u) if (bcast->type == BROADCAST_ROW) { /* (1, n) -> (m, n): replicate row m times */ - for (int j = 0; j < bcast->n; j++) + for (int j = 0; j < node->d2; j++) { - for (int i = 0; i < bcast->m; i++) + for (int i = 0; i < node->d1; i++) { - node->value[i + j * bcast->m] = x->value[j]; + node->value[i + j * node->d1] = x->value[j]; } } } else if (bcast->type == BROADCAST_COL) { /* (m, 1) -> (m, n): replicate column n times */ - for (int j = 0; j < bcast->n; j++) + for (int j = 0; j < node->d2; j++) { - memcpy(node->value + j * bcast->m, x->value, bcast->m * sizeof(double)); + memcpy(node->value + j * node->d1, x->value, node->d1 * sizeof(double)); } } else @@ -62,17 +62,17 @@ static void jacobian_init(expr *node) if (bcast->type == BROADCAST_ROW) { /* Row broadcast: (1, n) -> (m, n) */ - total_nnz = x->jacobian->nnz * bcast->m; + total_nnz = x->jacobian->nnz * node->d1; } else if (bcast->type == BROADCAST_COL) { /* Column broadcast: (m, 1) -> (m, n) */ - total_nnz = x->jacobian->nnz * bcast->n; + total_nnz = x->jacobian->nnz * node->d2; } else { /* Scalar broadcast: (1, 1) -> (m, n) */ - total_nnz = x->jacobian->nnz * bcast->m * bcast->n; + total_nnz = x->jacobian->nnz * node->d1 * node->d2; } node->jacobian = new_csr_matrix(node->size, node->n_vars, total_nnz); @@ -86,17 +86,17 @@ static void jacobian_init(expr *node) if (bcast->type == BROADCAST_ROW) { - for (int i = 0; i < bcast->n; i++) + for (int i = 0; i < node->d2; i++) { int nnz_in_row = Jx->p[i + 1] - Jx->p[i]; /* copy columns indices */ - tile_int(J->i + J->nnz, Jx->i + Jx->p[i], nnz_in_row, bcast->m); + tile_int(J->i + J->nnz, Jx->i + Jx->p[i], nnz_in_row, node->d1); /* set row pointers */ - for (int rep = 0; rep < bcast->m; rep++) + for (int rep = 0; rep < node->d1; rep++) { - J->p[i * bcast->m + rep] = J->nnz; + J->p[i * node->d1 + rep] = J->nnz; J->nnz += nnz_in_row; } } @@ -106,15 +106,15 @@ static void jacobian_init(expr *node) { /* copy column indices */ - tile_int(J->i, Jx->i, Jx->nnz, bcast->n); + tile_int(J->i, Jx->i, Jx->nnz, node->d2); /* set row pointers */ int offset = 0; - for (int i = 0; i < bcast->n; i++) + for (int i = 0; i < node->d2; i++) { - for (int j = 0; j < bcast->m; j++) + for (int j = 0; j < node->d1; j++) { - J->p[i * bcast->m + j] = offset; + J->p[i * node->d1 + j] = offset; offset += Jx->p[1] - Jx->p[0]; } } @@ -124,12 +124,12 @@ static void jacobian_init(expr *node) else { /* copy column indices */ - tile_int(J->i, Jx->i, Jx->nnz, bcast->m * bcast->n); + tile_int(J->i, Jx->i, Jx->nnz, node->d1 * node->d2); /* set row pointers */ int offset = 0; int nnz = Jx->p[1] - Jx->p[0]; - for (int i = 0; i < bcast->m * bcast->n; i++) + for (int i = 0; i < node->d1 * node->d2; i++) { J->p[i] = offset; offset += nnz; @@ -150,20 +150,20 @@ static void eval_jacobian(expr *node) if (bcast->type == BROADCAST_ROW) { - for (int i = 0; i < bcast->n; i++) + for (int i = 0; i < node->d2; i++) { int nnz_in_row = Jx->p[i + 1] - Jx->p[i]; - tile_double(J->x + J->nnz, Jx->x + Jx->p[i], nnz_in_row, bcast->m); - J->nnz += nnz_in_row * bcast->m; + tile_double(J->x + J->nnz, Jx->x + Jx->p[i], nnz_in_row, node->d1); + J->nnz += nnz_in_row * node->d1; } } else if (bcast->type == BROADCAST_COL) { - tile_double(J->x, Jx->x, Jx->nnz, bcast->n); + tile_double(J->x, Jx->x, Jx->nnz, node->d2); } else { - tile_double(J->x, Jx->x, Jx->nnz, bcast->m * bcast->n); + tile_double(J->x, Jx->x, Jx->nnz, node->d1 * node->d2); } } @@ -192,22 +192,22 @@ static void eval_wsum_hess(expr *node, const double *w) if (bcast->type == BROADCAST_ROW) { /* (1, n) -> (m, n): each input element has m weights to sum */ - for (int j = 0; j < bcast->n; j++) + for (int j = 0; j < node->d2; j++) { - for (int i = 0; i < bcast->m; i++) + for (int i = 0; i < node->d1; i++) { - node->dwork[j] += w[i + j * bcast->m]; + node->dwork[j] += w[i + j * node->d1]; } } } else if (bcast->type == BROADCAST_COL) { /* (m, 1) -> (m, n): each input element has n weights to sum */ - for (int j = 0; j < bcast->n; j++) + for (int j = 0; j < node->d2; j++) { - for (int i = 0; i < bcast->m; i++) + for (int i = 0; i < node->d1; i++) { - node->dwork[i] += w[i + j * bcast->m]; + node->dwork[i] += w[i + j * node->d1]; } } } @@ -215,7 +215,7 @@ static void eval_wsum_hess(expr *node, const double *w) { /* (1, 1) -> (m, n): scalar has m*n weights to sum */ node->dwork[0] = 0.0; - for (int k = 0; k < bcast->m * bcast->n; k++) + for (int k = 0; k < node->size; k++) { node->dwork[0] += w[k]; } @@ -230,20 +230,18 @@ static bool is_affine(const expr *node) return node->left->is_affine(node->left); } -expr *new_broadcast(expr *child, int target_d1, int target_d2) +expr *new_broadcast(expr *child, int d1, int d2) { // --------------------------------------------------------------------------- // determine broadcast type // --------------------------------------------------------------------------- broadcast_type type; - int m = target_d1; - int n = target_d2; - if (child->d1 == 1 && child->d2 == n) + if (child->d1 == 1 && child->d2 == d2) { type = BROADCAST_ROW; } - else if (child->d1 == m && child->d2 == 1) + else if (child->d1 == d1 && child->d2 == 1) { type = BROADCAST_COL; } @@ -265,15 +263,13 @@ expr *new_broadcast(expr *child, int target_d1, int target_d2) // -------------------------------------------------------------------------- // initialize the rest of the expression // -------------------------------------------------------------------------- - init_expr(node, target_d1, target_d2, child->n_vars, forward, jacobian_init, - eval_jacobian, is_affine, NULL); + init_expr(node, d1, d2, child->n_vars, forward, jacobian_init, eval_jacobian, + is_affine, NULL); node->left = child; expr_retain(child); node->wsum_hess_init = wsum_hess_init; node->eval_wsum_hess = eval_wsum_hess; bcast->type = type; - bcast->m = m; - bcast->n = n; return node; } diff --git a/src/affine/index.c b/src/affine/index.c index 49948a9..d88b990 100644 --- a/src/affine/index.c +++ b/src/affine/index.c @@ -74,9 +74,8 @@ static void eval_jacobian(expr *node) for (int i = 0; i < idx->n_idxs; i++) { - int row = idx->indices[i]; int len = J->p[i + 1] - J->p[i]; - memcpy(J->x + J->p[i], Jx->x + Jx->p[row], len * sizeof(double)); + memcpy(J->x + J->p[i], Jx->x + Jx->p[idx->indices[i]], len * sizeof(double)); } } diff --git a/src/affine/neg.c b/src/affine/neg.c index 71be271..c657aff 100644 --- a/src/affine/neg.c +++ b/src/affine/neg.c @@ -16,17 +16,16 @@ static void forward(expr *node, const double *u) static void jacobian_init(expr *node) { + expr *x = node->left; /* initialize child's jacobian */ - node->left->jacobian_init(node->left); + x->jacobian_init(x); /* same sparsity pattern as child */ - CSR_Matrix *child_jac = node->left->jacobian; - node->jacobian = new_csr_matrix(child_jac->m, child_jac->n, child_jac->nnz); + node->jacobian = new_csr_matrix(node->size, node->n_vars, x->jacobian->nnz); /* copy row pointers and column indices (sparsity pattern is constant) */ - memcpy(node->jacobian->p, child_jac->p, (child_jac->m + 1) * sizeof(int)); - memcpy(node->jacobian->i, child_jac->i, child_jac->nnz * sizeof(int)); - node->jacobian->nnz = child_jac->nnz; + memcpy(node->jacobian->p, x->jacobian->p, (x->jacobian->m + 1) * sizeof(int)); + memcpy(node->jacobian->i, x->jacobian->i, x->jacobian->nnz * sizeof(int)); } static void eval_jacobian(expr *node) diff --git a/src/affine/reshape.c b/src/affine/reshape.c index d5aeeba..ef89947 100644 --- a/src/affine/reshape.c +++ b/src/affine/reshape.c @@ -32,20 +32,19 @@ static void eval_jacobian(expr *node) static void wsum_hess_init(expr *node) { - node->left->wsum_hess_init(node->left); - CSR_Matrix *child_hess = node->left->wsum_hess; - node->wsum_hess = new_csr_matrix(child_hess->m, child_hess->n, child_hess->nnz); - memcpy(node->wsum_hess->p, child_hess->p, (child_hess->m + 1) * sizeof(int)); - memcpy(node->wsum_hess->i, child_hess->i, child_hess->nnz * sizeof(int)); - node->wsum_hess->nnz = child_hess->nnz; + expr *x = node->left; + x->wsum_hess_init(x); + node->wsum_hess = + new_csr_matrix(x->wsum_hess->m, x->wsum_hess->n, x->wsum_hess->nnz); + memcpy(node->wsum_hess->p, x->wsum_hess->p, (x->wsum_hess->m + 1) * sizeof(int)); + memcpy(node->wsum_hess->i, x->wsum_hess->i, x->wsum_hess->nnz * sizeof(int)); } static void eval_wsum_hess(expr *node, const double *w) { - node->left->eval_wsum_hess(node->left, w); - CSR_Matrix *child_hess = node->left->wsum_hess; - CSR_Matrix *hess = node->wsum_hess; - memcpy(hess->x, child_hess->x, child_hess->nnz * sizeof(double)); + expr *x = node->left; + x->eval_wsum_hess(x, w); + memcpy(node->wsum_hess->x, x->wsum_hess->x, x->wsum_hess->nnz * sizeof(double)); } static bool is_affine(const expr *node) diff --git a/tests/jacobian_tests/test_index.h b/tests/jacobian_tests/test_index.h index 00b7cb2..5c156a0 100644 --- a/tests/jacobian_tests/test_index.h +++ b/tests/jacobian_tests/test_index.h @@ -14,7 +14,7 @@ const char *test_index_forward_simple(void) double u[3] = {1.0, 2.0, 3.0}; int indices[2] = {0, 2}; expr *var = new_variable(3, 1, 0, 3); - expr *idx = new_index(var, indices, 2); + expr *idx = new_index(var, 1, 2, indices, 2); idx->forward(idx, u); double expected[2] = {1.0, 3.0}; @@ -30,7 +30,7 @@ const char *test_index_forward_repeated(void) double u[3] = {1.0, 2.0, 3.0}; int indices[3] = {0, 0, 2}; expr *var = new_variable(3, 1, 0, 3); - expr *idx = new_index(var, indices, 3); + expr *idx = new_index(var, 1, 3, indices, 3); idx->forward(idx, u); double expected[3] = {1.0, 1.0, 3.0}; @@ -47,7 +47,7 @@ const char *test_index_jacobian_of_variable(void) double u[3] = {1.0, 2.0, 3.0}; int indices[2] = {0, 2}; expr *var = new_variable(3, 1, 0, 3); - expr *idx = new_index(var, indices, 2); + expr *idx = new_index(var, 1, 2, indices, 2); idx->forward(idx, u); idx->jacobian_init(idx); idx->eval_jacobian(idx); @@ -72,7 +72,7 @@ const char *test_index_jacobian_of_log(void) int indices[2] = {0, 2}; expr *var = new_variable(3, 1, 0, 3); expr *log_node = new_log(var); - expr *idx = new_index(log_node, indices, 2); + expr *idx = new_index(log_node, 1, 2, indices, 2); idx->forward(idx, u); idx->jacobian_init(idx); idx->eval_jacobian(idx); @@ -98,7 +98,7 @@ const char *test_index_jacobian_repeated(void) double u[3] = {1.0, 2.0, 3.0}; int indices[2] = {0, 0}; expr *var = new_variable(3, 1, 0, 3); - expr *idx = new_index(var, indices, 2); + expr *idx = new_index(var, 1, 2, indices, 2); idx->forward(idx, u); idx->jacobian_init(idx); idx->eval_jacobian(idx); @@ -127,7 +127,7 @@ const char *test_sum_of_index(void) int indices[2] = {0, 2}; expr *var = new_variable(3, 1, 0, 3); - expr *idx = new_index(var, indices, 2); + expr *idx = new_index(var, 1, 2, indices, 2); expr *s = new_sum(idx, -1); /* sum all */ s->forward(s, u); diff --git a/tests/wsum_hess/test_index.h b/tests/wsum_hess/test_index.h index 0a2fb9c..4f78362 100644 --- a/tests/wsum_hess/test_index.h +++ b/tests/wsum_hess/test_index.h @@ -23,7 +23,7 @@ const char *test_wsum_hess_index_log(void) expr *var = new_variable(3, 1, 0, 3); expr *log_node = new_log(var); - expr *idx = new_index(log_node, indices, 2); + expr *idx = new_index(log_node, 1, 2, indices, 2); idx->forward(idx, u); idx->jacobian_init(idx); @@ -38,7 +38,8 @@ const char *test_wsum_hess_index_log(void) int expected_p[4] = {0, 1, 2, 3}; int expected_i[3] = {0, 1, 2}; - mu_assert("index log hess vals", cmp_double_array(idx->wsum_hess->x, expected_x, 3)); + mu_assert("index log hess vals", + cmp_double_array(idx->wsum_hess->x, expected_x, 3)); mu_assert("index log hess p", cmp_int_array(idx->wsum_hess->p, expected_p, 4)); mu_assert("index log hess i", cmp_int_array(idx->wsum_hess->i, expected_i, 3)); @@ -56,7 +57,7 @@ const char *test_wsum_hess_index_repeated(void) expr *var = new_variable(3, 1, 0, 3); expr *log_node = new_log(var); - expr *idx = new_index(log_node, indices, 2); + expr *idx = new_index(log_node, 1, 2, indices, 2); idx->forward(idx, u); idx->jacobian_init(idx); @@ -70,9 +71,12 @@ const char *test_wsum_hess_index_repeated(void) int expected_p[4] = {0, 1, 2, 3}; int expected_i[3] = {0, 1, 2}; - mu_assert("index repeated hess vals", cmp_double_array(idx->wsum_hess->x, expected_x, 3)); - mu_assert("index repeated hess p", cmp_int_array(idx->wsum_hess->p, expected_p, 4)); - mu_assert("index repeated hess i", cmp_int_array(idx->wsum_hess->i, expected_i, 3)); + mu_assert("index repeated hess vals", + cmp_double_array(idx->wsum_hess->x, expected_x, 3)); + mu_assert("index repeated hess p", + cmp_int_array(idx->wsum_hess->p, expected_p, 4)); + mu_assert("index repeated hess i", + cmp_int_array(idx->wsum_hess->i, expected_i, 3)); free_expr(idx); return 0; @@ -90,7 +94,7 @@ const char *test_wsum_hess_sum_index_log(void) expr *var = new_variable(3, 1, 0, 3); expr *log_node = new_log(var); - expr *idx = new_index(log_node, indices, 2); + expr *idx = new_index(log_node, 1, 2, indices, 2); expr *sum_node = new_sum(idx, -1); sum_node->forward(sum_node, u); @@ -108,8 +112,10 @@ const char *test_wsum_hess_sum_index_log(void) mu_assert("sum index log hess vals", cmp_double_array(sum_node->wsum_hess->x, expected_x, 3)); - mu_assert("sum index log hess p", cmp_int_array(sum_node->wsum_hess->p, expected_p, 4)); - mu_assert("sum index log hess i", cmp_int_array(sum_node->wsum_hess->i, expected_i, 3)); + mu_assert("sum index log hess p", + cmp_int_array(sum_node->wsum_hess->p, expected_p, 4)); + mu_assert("sum index log hess i", + cmp_int_array(sum_node->wsum_hess->i, expected_i, 3)); free_expr(sum_node); return 0; From b30ad5eedc6f0c05a730fec9e299d1ec7a091e09 Mon Sep 17 00:00:00 2001 From: Daniel Date: Sat, 17 Jan 2026 04:37:52 -0800 Subject: [PATCH 7/8] left matmul, internal logic for handling numpy's weird broadcasting rule --- src/bivariate/left_matmul.c | 32 +++++++++++++++++++++++++++++--- 1 file changed, 29 insertions(+), 3 deletions(-) diff --git a/src/bivariate/left_matmul.c b/src/bivariate/left_matmul.c index dba2097..9a72c42 100644 --- a/src/bivariate/left_matmul.c +++ b/src/bivariate/left_matmul.c @@ -1,6 +1,7 @@ #include "bivariate.h" #include "subexpr.h" #include +#include #include /* This file implement the atom 'left_matmul' corresponding to the operation y = @@ -114,12 +115,37 @@ static void eval_wsum_hess(expr *node, const double *w) expr *new_left_matmul(expr *u, const CSR_Matrix *A) { - assert(u->d1 == A->n); + /* We expect u->d1 == A->n. However, numpy's broadcasting rules allow users to do + A @ u where u is (n, ) which in C is actually (1, n). In that case the result + of A @ u is (m, ), which is (1, m) according to broadcasting rules. We + therefore check if this is the case. */ + int d1, d2, n_blocks; + if (u->d1 == A->n) + { + d1 = A->m; + d2 = u->d2; + n_blocks = u->d2; + printf("d2: %d\n", d2); + } + else if (u->d2 == A->n && u->d1 == 1) + { + d1 = 1; + d2 = A->m; + n_blocks = 1; + printf("d2: %d\n", d2); + } + else + { + fprintf(stderr, "Error in new_left_matmul: dimension mismatch \n"); + exit(1); + } + + // assert(u->d1 == A->n); /* Allocate the type-specific struct */ left_matmul_expr *lin_node = (left_matmul_expr *) calloc(1, sizeof(left_matmul_expr)); expr *node = &lin_node->base; - init_expr(node, A->m, u->d2, u->n_vars, forward, jacobian_init, eval_jacobian, + init_expr(node, d1, d2, u->n_vars, forward, jacobian_init, eval_jacobian, is_affine, free_type_data); node->left = u; expr_retain(u); @@ -127,7 +153,7 @@ expr *new_left_matmul(expr *u, const CSR_Matrix *A) node->eval_wsum_hess = eval_wsum_hess; /* Initialize type-specific fields */ - lin_node->A = block_diag_repeat_csr(A, node->d2); + lin_node->A = block_diag_repeat_csr(A, n_blocks); int alloc = MAX(lin_node->A->n, node->n_vars); node->iwork = (int *) malloc(alloc * sizeof(int)); lin_node->AT = transpose(lin_node->A, node->iwork); From 0a7b8a66b4dffff42ab6feee88f62361927a7182 Mon Sep 17 00:00:00 2001 From: Daniel Date: Sat, 17 Jan 2026 04:45:28 -0800 Subject: [PATCH 8/8] minor cleanup --- TODO.md | 4 ++-- src/bivariate/left_matmul.c | 3 --- src/bivariate/quad_over_lin.c | 3 --- 3 files changed, 2 insertions(+), 8 deletions(-) diff --git a/TODO.md b/TODO.md index ec3ed93..64dc697 100644 --- a/TODO.md +++ b/TODO.md @@ -1,10 +1,10 @@ 3. more tests for chain rule elementwise univariate hessian 4. in the refactor, add consts -5. multiply with one constant vector/scalar argument 6. AX where X is a matrix. Can that happen? How is that canonicalized? Maybe it can't happen. 7. Must be able to compute jacobian and hessian of A @ phi(x), so linear operator needs other code! This requires new infrastructure, I think. 8. Shortcut hessians of affine stuff? -9. quad-over-lin with first argument matrix. +10. For performance reasons, is it useful to have a dense matmul with A and B as dense matrices? +11. right matmul, add broadcasting logic as in left matmul Going through all atoms to see that sparsity pattern is computed in initialization of jacobian: 2. trace - not ok diff --git a/src/bivariate/left_matmul.c b/src/bivariate/left_matmul.c index 9a72c42..0ead275 100644 --- a/src/bivariate/left_matmul.c +++ b/src/bivariate/left_matmul.c @@ -125,14 +125,12 @@ expr *new_left_matmul(expr *u, const CSR_Matrix *A) d1 = A->m; d2 = u->d2; n_blocks = u->d2; - printf("d2: %d\n", d2); } else if (u->d2 == A->n && u->d1 == 1) { d1 = 1; d2 = A->m; n_blocks = 1; - printf("d2: %d\n", d2); } else { @@ -140,7 +138,6 @@ expr *new_left_matmul(expr *u, const CSR_Matrix *A) exit(1); } - // assert(u->d1 == A->n); /* Allocate the type-specific struct */ left_matmul_expr *lin_node = (left_matmul_expr *) calloc(1, sizeof(left_matmul_expr)); diff --git a/src/bivariate/quad_over_lin.c b/src/bivariate/quad_over_lin.c index 47f5714..ce28685 100644 --- a/src/bivariate/quad_over_lin.c +++ b/src/bivariate/quad_over_lin.c @@ -292,9 +292,6 @@ static void eval_wsum_hess(expr *node, const double *w) } } -/* TODO: the above implementations probably assumes that the first argument is a - vector but cvxpy supports a matrix for the first argument*/ - expr *new_quad_over_lin(expr *left, expr *right) { assert((right->d2 == 1 && right->d2 == 1)); /* right must be scalar*/