Add native kron (Kronecker product) affine atom

include/atoms/affine.h

@@ -80,4 +80,18 @@ expr *new_vector_mult(expr *param_node, expr *child);
    kernel and may either represent a constant or an updatable parameter */
 expr *new_convolve(expr *param_node, expr *child);
 
+/* Kronecker product Z = kron(A, B), built sparse-only: one operand is the
+   variable-free constant/parameter (param_node), the other carries the variables
+   (child). active_blocks holds the column-major indices of the constant operand's
+   active (nonzero) blocks -- only those output rows are materialized. For a
+   parametric operand cvxpy passes all blocks (dense, slow). (p, q) are A's dims,
+   (r, s) are B's dims.
+
+   left_kron:  A = param_node (p x q), B = child (r x s); active_blocks index A.
+   right_kron: A = child (p x q), B = param_node (r x s); active_blocks index B. */
+expr *new_left_kron(expr *param_node, expr *child, int p, int q, int r, int s,
+                    const int *active_blocks, int n_active);
+expr *new_right_kron(expr *param_node, expr *child, int p, int q, int r, int s,
+                     const int *active_blocks, int n_active);
+
 #endif /* AFFINE_H */

include/subexpr.h

@@ -173,6 +173,27 @@ typedef struct convolve_expr
     CSC_matrix *Jchild_CSC;
 } convolve_expr;
 
+/* Kronecker product Z = kron(A, B) where one operand is variable-free (the
+ * constant/parameter, held by param_source) and the other (child = node->left)
+ * carries the variables. Built sparse-only by new_left_kron / new_right_kron:
+ * cvxpy passes the constant operand's active (nonzero) block indices, so only
+ * those output rows are filled; the rest stay inactive (child_row == -1) and
+ * contribute a zero value and an empty Jacobian row. Every active output entry
+ * depends on a single child entry, so the Jacobian is the child Jacobian's rows
+ * gathered (with repetition) and scaled by the constant -- no coefficient matrix
+ * or matmul. */
+typedef struct kron_expr
+{
+    expr base;
+    expr *param_source; /* the constant/parameter operand node (re-evaluated each
+                           solve) */
+    int p, q, r, s;     /* A is p x q, B is r x s */
+    int *child_row; /* size_out: child entry each output row gathers; -1 if inactive
+                     */
+    int *coeff_idx; /* size_out: index into param_source->value (valid where
+                       child_row >= 0) */
+} kron_expr;
+
 /* Bivariate matrix multiplication: Z = f(u) @ g(u) where both children
  * may be composite expressions. */
 typedef struct matmul_expr

src/atoms/affine/kron.c

@@ -0,0 +1,271 @@
+/*
+ * Copyright 2026 Daniel Cederberg and William Zhang
+ *
+ * This file is part of the SparseDiffEngine project.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#include "atoms/affine.h"
+#include "subexpr.h"
+#include "utils/CSR_matrix.h"
+#include "utils/sparse_matrix.h"
+#include "utils/tracked_alloc.h"
+#include <assert.h>
+#include <stdio.h>
+#include <stdlib.h>
+#include <string.h>
+
+/* Kronecker product Z = kron(A, B), where one operand is variable-free
+ * (param_source) and the other (child = node->left) carries the variables.
+ *
+ * Built sparse-only: cvxpy passes the constant operand's active (nonzero) block
+ * indices, and new_left_kron / new_right_kron fill child_row[]/coeff_idx[] only
+ * for the output rows those blocks cover. Inactive output rows keep
+ * child_row == -1 (a zero kron entry) and contribute an empty Jacobian row.
+ *
+ * With column-major (Fortran) flattening, an output index OUT = I + J*(p*r)
+ * decomposes as I = i*r + k and J = j*s + l. The output block (i, j) inner
+ * (k, l) equals A[i,j] * B[k,l], so every active output entry depends on a
+ * single child entry:
+ *
+ *   Z[OUT]        = coeff[OUT] * vec(child)[child_row[OUT]]
+ *   J_kron[OUT,:] = coeff[OUT] * J_child[child_row[OUT], :]
+ *
+ * where coeff[OUT] = param_source->value[coeff_idx[OUT]]. The shared vtable
+ * below honors the child_row == -1 sentinel; only the construction loop differs
+ * between left_kron and right_kron. */
+
+/* Pull current parameter values through any broadcast/promote wrappers. */
+static void refresh_param_values(kron_expr *knode)
+{
+    if (!knode->base.needs_parameter_refresh)
+    {
+        return;
+    }
+
+    knode->param_source->forward(knode->param_source, NULL);
+    knode->base.needs_parameter_refresh = false;
+}
+
+static void forward(expr *node, const double *u)
+{
+    expr *child = node->left;
+    kron_expr *knode = (kron_expr *) node;
+
+    refresh_param_values(knode);
+    child->forward(child, u);
+
+    const double *a = knode->param_source->value;
+    const double *x = child->value;
+    double *y = node->value;
+    for (int out = 0; out < node->size; out++)
+    {
+        int cr = knode->child_row[out];
+        y[out] = (cr < 0) ? 0.0 : a[knode->coeff_idx[out]] * x[cr];
+    }
+}
+
+static void jacobian_init_impl(expr *node)
+{
+    expr *child = node->left;
+    kron_expr *knode = (kron_expr *) node;
+
+    jacobian_init(child);
+
+    /* Active output row OUT shares the column set of child row child_row[OUT];
+       inactive rows (child_row == -1) are empty. Build the result CSR sparsity
+       by copying the active child rows (with repetition). */
+    CSR_matrix *Jc = child->jacobian->to_csr(child->jacobian);
+
+    int total = 0;
+    for (int out = 0; out < node->size; out++)
+    {
+        int cr = knode->child_row[out];
+        if (cr >= 0) total += Jc->p[cr + 1] - Jc->p[cr];
+    }
+
+    CSR_matrix *Jk = new_CSR_matrix(node->size, node->n_vars, total);
+    int idx = 0;
+    Jk->p[0] = 0;
+    for (int out = 0; out < node->size; out++)
+    {
+        int cr = knode->child_row[out];
+        if (cr >= 0)
+        {
+            for (int t = Jc->p[cr]; t < Jc->p[cr + 1]; t++)
+            {
+                Jk->i[idx++] = Jc->i[t];
+            }
+        }
+        Jk->p[out + 1] = idx;
+    }
+    node->jacobian = new_sparse_matrix(Jk);
+}
+
+static void eval_jacobian(expr *node)
+{
+    expr *child = node->left;
+    kron_expr *knode = (kron_expr *) node;
+
+    child->eval_jacobian(child);
+
+    /* Child sparsity is fixed after jacobian_init, so the result row offsets
+       still align; refill active rows as scale * child-row-values. */
+    CSR_matrix *Jc = child->jacobian->to_csr(child->jacobian);
+    CSR_matrix *Jk = node->jacobian->to_csr(node->jacobian);
+    const double *a = knode->param_source->value;
+
+    int idx = 0;
+    for (int out = 0; out < node->size; out++)
+    {
+        int cr = knode->child_row[out];
+        if (cr < 0) continue;
+        double scale = a[knode->coeff_idx[out]];
+        for (int t = Jc->p[cr]; t < Jc->p[cr + 1]; t++)
+        {
+            Jk->x[idx++] = scale * Jc->x[t];
+        }
+    }
+}
+
+static void wsum_hess_init_impl(expr *node)
+{
+    expr *child = node->left;
+
+    wsum_hess_init(child);
+    node->wsum_hess = child->wsum_hess->copy_sparsity(child->wsum_hess);
+    /* backprop workspace: one weight per child entry */
+    node->work->dwork = (double *) sp_malloc(child->size * sizeof(double));
+}
+
+static void eval_wsum_hess(expr *node, const double *w)
+{
+    expr *child = node->left;
+    kron_expr *knode = (kron_expr *) node;
+    const double *a = knode->param_source->value;
+    double *w_prime = node->work->dwork;
+
+    /* kron is affine in child, so the Hessian is the child's with weights pushed
+       back through the linear gather: w'[child_row] += coeff * w[OUT]. Many
+       output rows map to one child entry, hence the accumulation. */
+    memset(w_prime, 0, child->size * sizeof(double));
+    for (int out = 0; out < node->size; out++)
+    {
+        int cr = knode->child_row[out];
+        if (cr >= 0) w_prime[cr] += a[knode->coeff_idx[out]] * w[out];
+    }
+
+    child->eval_wsum_hess(child, w_prime);
+    memcpy(node->wsum_hess->x, child->wsum_hess->x,
+           node->wsum_hess->nnz * sizeof(double));
+}
+
+static bool is_affine(const expr *node)
+{
+    return node->left->is_affine(node->left);
+}
+
+static void free_type_data(expr *node)
+{
+    kron_expr *knode = (kron_expr *) node;
+    sp_free(knode->child_row);
+    sp_free(knode->coeff_idx);
+    free_expr(knode->param_source);
+
+    knode->child_row = NULL;
+    knode->coeff_idx = NULL;
+    knode->param_source = NULL;
+}
+
+/* Allocate a kron node and its (all-inactive) index arrays. The left/right
+   constructors then fill the active rows. */
+static kron_expr *new_kron_common(expr *param_node, expr *child, int p, int q, int r,
+                                  int s)
+{
+    int size_out = (p * r) * (q * s);
+
+    kron_expr *knode = (kron_expr *) sp_calloc(1, sizeof(kron_expr));
+    expr *node = &knode->base;
+    init_expr(node, p * r, q * s, child->n_vars, forward, jacobian_init_impl,
+              eval_jacobian, is_affine, wsum_hess_init_impl, eval_wsum_hess,
+              free_type_data);
+    node->left = child;
+    expr_retain(child);
+
+    knode->param_source = param_node;
+    expr_retain(param_node);
+    knode->p = p;
+    knode->q = q;
+    knode->r = r;
+    knode->s = s;
+
+    knode->child_row = (int *) sp_malloc(size_out * sizeof(int));
+    knode->coeff_idx = (int *) sp_malloc(size_out * sizeof(int));
+    for (int out = 0; out < size_out; out++)
+        knode->child_row[out] = -1; /* inactive until an active block fills it */
+
+    knode->base.needs_parameter_refresh = true;
+    return knode;
+}
+
+/* Z = kron(A, B) with A = param_node (p x q) the constant, B = child (r x s) the
+   variable. active_blocks holds column-major indices i + j*p of A's nonzeros. */
+expr *new_left_kron(expr *param_node, expr *child, int p, int q, int r, int s,
+                    const int *active_blocks, int n_active)
+{
+    kron_expr *knode = new_kron_common(param_node, child, p, q, r, s);
+    int n_rows = p * r;
+    for (int b = 0; b < n_active; b++)
+    {
+        int bidx = active_blocks[b]; /* = i + j*p into A */
+        assert(0 <= bidx && bidx < p * q);
+        int i = bidx % p;
+        int j = bidx / p;
+        for (int l = 0; l < s; l++)
+        {
+            for (int k = 0; k < r; k++)
+            {
+                int out = (i * r + k) + (j * s + l) * n_rows;
+                knode->child_row[out] = k + l * r; /* col-major into B */
+                knode->coeff_idx[out] = bidx;      /* col-major into A */
+            }
+        }
+    }
+    return &knode->base;
+}
+
+/* Z = kron(A, B) with A = child (p x q) the variable, B = param_node (r x s) the
+   constant. active_blocks holds column-major indices k + l*r of B's nonzeros. */
+expr *new_right_kron(expr *param_node, expr *child, int p, int q, int r, int s,
+                     const int *active_blocks, int n_active)
+{
+    kron_expr *knode = new_kron_common(param_node, child, p, q, r, s);
+    int n_rows = p * r;
+    for (int b = 0; b < n_active; b++)
+    {
+        int bidx = active_blocks[b]; /* = k + l*r into B */
+        assert(0 <= bidx && bidx < r * s);
+        int k = bidx % r;
+        int l = bidx / r;
+        for (int j = 0; j < q; j++)
+        {
+            for (int i = 0; i < p; i++)
+            {
+                int out = (i * r + k) + (j * s + l) * n_rows;
+                knode->child_row[out] = i + j * p; /* col-major into A */
+                knode->coeff_idx[out] = bidx;      /* col-major into B */
+            }
+        }
+    }
+    return &knode->base;
+}

tests/all_tests.c

@@ -10,6 +10,7 @@
 #include "forward_pass/affine/test_convolve.h"
 #include "forward_pass/affine/test_diag_mat.h"
 #include "forward_pass/affine/test_hstack.h"
+#include "forward_pass/affine/test_kron.h"
 #include "forward_pass/affine/test_left_matmul_dense.h"
 #include "forward_pass/affine/test_linear_op.h"
 #include "forward_pass/affine/test_neg.h"
@@ -30,6 +31,7 @@
 #include "jacobian_tests/affine/test_diag_mat.h"
 #include "jacobian_tests/affine/test_hstack.h"
 #include "jacobian_tests/affine/test_index.h"
+#include "jacobian_tests/affine/test_kron.h"
 #include "jacobian_tests/affine/test_left_matmul.h"
 #include "jacobian_tests/affine/test_neg.h"
 #include "jacobian_tests/affine/test_promote.h"
@@ -76,6 +78,7 @@
 #include "wsum_hess/affine/test_diag_mat.h"
 #include "wsum_hess/affine/test_hstack.h"
 #include "wsum_hess/affine/test_index.h"
+#include "wsum_hess/affine/test_kron.h"
 #include "wsum_hess/affine/test_left_matmul.h"
 #include "wsum_hess/affine/test_right_matmul.h"
 #include "wsum_hess/affine/test_scalar_mult.h"
@@ -151,6 +154,10 @@ int main(void)
     mu_run_test(test_convolve_forward, tests_run);
     mu_run_test(test_convolve_forward_row, tests_run);
     mu_run_test(test_convolve_forward_param, tests_run);
+    mu_run_test(test_kron_forward_const_left, tests_run);
+    mu_run_test(test_kron_forward_const_right, tests_run);
+    mu_run_test(test_kron_forward_scalar, tests_run);
+    mu_run_test(test_kron_forward_sparse, tests_run);
     mu_run_test(test_diag_mat_forward, tests_run);
     mu_run_test(test_upper_tri_forward_4x4, tests_run);
 
@@ -244,6 +251,10 @@ int main(void)
     mu_run_test(test_jacobian_matmul, tests_run);
     mu_run_test(test_jacobian_convolve, tests_run);
     mu_run_test(test_jacobian_convolve_composite, tests_run);
+    mu_run_test(test_jacobian_kron_const_left, tests_run);
+    mu_run_test(test_jacobian_kron_const_right, tests_run);
+    mu_run_test(test_jacobian_kron_sparse, tests_run);
+    mu_run_test(test_jacobian_kron_composite, tests_run);
     mu_run_test(test_jacobian_transpose, tests_run);
     mu_run_test(test_jacobian_transpose_pd_preserved, tests_run);
     mu_run_test(test_diag_mat_jacobian_variable, tests_run);
@@ -317,6 +328,8 @@ int main(void)
     mu_run_test(test_wsum_hess_right_matmul_vector, tests_run);
     mu_run_test(test_wsum_hess_convolve, tests_run);
     mu_run_test(test_wsum_hess_convolve_composite, tests_run);
+    mu_run_test(test_wsum_hess_kron, tests_run);
+    mu_run_test(test_wsum_hess_kron_composite, tests_run);
     mu_run_test(test_wsum_hess_broadcast_row, tests_run);
     mu_run_test(test_wsum_hess_broadcast_col, tests_run);
     mu_run_test(test_wsum_hess_broadcast_scalar_to_matrix, tests_run);