[WIP] Add kron_left bivariate atom

include/bivariate.h

@@ -42,6 +42,10 @@ expr *new_right_matmul(expr *u, const CSR_Matrix *A);
 
 expr *new_right_matmul_dense(expr *u, int m, int n, const double *data);
 
+/* Kronecker product: kron(C, X) where C is a constant sparse matrix and X is
+ * an expression of shape (p x q). Output has shape (m*p, n*q). */
+expr *new_kron_left(expr *child, const CSR_Matrix *C, int p, int q);
+
 /* Constant scalar multiplication: a * f(x) where a is a constant double */
 expr *new_const_scalar_mult(double a, expr *child);
 

include/subexpr.h

@@ -116,6 +116,16 @@ typedef struct left_matmul_expr
     int *csc_to_csr_work;
 } left_matmul_expr;
 
+/* Kronecker product: Z = kron(C, X) where C is a constant matrix */
+typedef struct kron_left_expr
+{
+    expr base;
+    int m, n, p, q;     /* C is (m x n), child X is (p x q) */
+    CSR_Matrix *C;      /* constant matrix, stored as CSR */
+    int *row_map;       /* output row -> child Jacobian row */
+    int *row_scale_idx; /* output row -> index into C->x for the scale factor */
+} kron_left_expr;
+
 /* Right matrix multiplication: y = f(x) * A where f(x) is an expression.
  * f(x) has shape p x n, A has shape n x q, output y has shape p x q.
  * Uses vec(y) = B * vec(f(x)) where B = A^T kron I_p. */

src/bivariate/kron_left.c

@@ -0,0 +1,338 @@
+/*
+ * Copyright 2026 Daniel Cederberg and William Zhang
+ *
+ * This file is part of the DNLP-differentiation-engine project.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#include "bivariate.h"
+#include "subexpr.h"
+#include <assert.h>
+#include <stdio.h>
+#include <stdlib.h>
+#include <string.h>
+
+/*
+ * Kronecker product: Z = kron(C, X) where C is a constant (m x n) matrix
+ * and X is a variable expression of shape (p x q).
+ *
+ * Output Z has shape (m*p, n*q), stored column-major as vec(Z) of length m*p*n*q.
+ *
+ * Key identity: Z[i*p+k, j*q+l] = C[i,j] * X[k,l]
+ * In column-major: vec(Z)[r] where r = (j*q+l)*(m*p) + i*p + k
+ *   depends on vec(X)[s] where s = l*p + k, with coefficient C[i,j].
+ *
+ * Jacobian structure: each row r of J_Z is a scaled copy of row s of J_X.
+ * Only rows where C[i,j] != 0 are non-trivial.
+ */
+
+/* ------------------------------------------------------------------ */
+/*                          Forward pass                               */
+/* ------------------------------------------------------------------ */
+static void forward(expr *node, const double *u)
+{
+    kron_left_expr *kn = (kron_left_expr *) node;
+    expr *child = node->left;
+    CSR_Matrix *C = kn->C;
+    int m = kn->m, p = kn->p, q = kn->q;
+    int mp = m * p;
+
+    child->forward(child, u);
+
+    /* Zero output first */
+    memset(node->value, 0, (size_t) node->size * sizeof(double));
+
+    /* For each nonzero C[i,j], fill block: Z[i*p+k, j*q+l] = C[i,j] * X[k,l] */
+    for (int i = 0; i < m; i++)
+    {
+        for (int idx = C->p[i]; idx < C->p[i + 1]; idx++)
+        {
+            int j = C->i[idx];
+            double cij = C->x[idx];
+
+            for (int l = 0; l < q; l++)
+            {
+                int z_col_start = (j * q + l) * mp + i * p;
+                int x_col_start = l * p;
+                for (int k = 0; k < p; k++)
+                {
+                    node->value[z_col_start + k] =
+                        cij * child->value[x_col_start + k];
+                }
+            }
+        }
+    }
+}
+
+/* ------------------------------------------------------------------ */
+/*                          Affine check                               */
+/* ------------------------------------------------------------------ */
+static bool is_affine(const expr *node)
+{
+    return node->left->is_affine(node->left);
+}
+
+/* ------------------------------------------------------------------ */
+/*                      Jacobian initialization                        */
+/* ------------------------------------------------------------------ */
+static void jacobian_init(expr *node)
+{
+    kron_left_expr *kn = (kron_left_expr *) node;
+    expr *child = node->left;
+    CSR_Matrix *C = kn->C;
+    int m = kn->m, p = kn->p, q = kn->q;
+    int mp = m * p;
+    int out_size = node->size; /* m * p * n * q */
+
+    /* Initialize child's Jacobian */
+    child->jacobian_init(child);
+    CSR_Matrix *Jchild = child->jacobian;
+
+    /*
+     * Count total nnz: for each output row r corresponding to nonzero C[i,j],
+     * the nnz equals the nnz of child's Jacobian row s = l*p + k.
+     */
+    int total_nnz = 0;
+    for (int i = 0; i < m; i++)
+    {
+        int row_nnz_C = C->p[i + 1] - C->p[i];
+        if (row_nnz_C == 0) continue;
+
+        for (int l = 0; l < q; l++)
+        {
+            for (int k = 0; k < p; k++)
+            {
+                int s = l * p + k; /* child row */
+                int child_row_nnz = Jchild->p[s + 1] - Jchild->p[s];
+                total_nnz += row_nnz_C * child_row_nnz;
+            }
+        }
+    }
+
+    /* Allocate Jacobian */
+    node->jacobian = new_csr_matrix(out_size, node->n_vars, total_nnz);
+
+    /* Allocate row_map and row_scale arrays for eval_jacobian */
+    kn->row_map = (int *) malloc((size_t) out_size * sizeof(int));
+    kn->row_scale_idx = (int *) malloc((size_t) out_size * sizeof(int));
+
+    /*
+     * Fill Jacobian sparsity pattern.
+     * Iterate in column-major order over Z:
+     *   r = (j*q + l) * mp + i*p + k
+     *   maps to child row s = l*p + k, with scale C[i,j]
+     */
+    int nnz_idx = 0;
+    for (int r = 0; r < out_size; r++)
+    {
+        node->jacobian->p[r] = nnz_idx;
+
+        /* Decode r: r = col_Z * mp + row_Z */
+        int row_Z = r % mp; /* i*p + k */
+        int col_Z = r / mp; /* j*q + l */
+        int i = row_Z / p;
+        int k = row_Z % p;
+        int j = col_Z / q;
+        int l = col_Z % q;
+        int s = l * p + k; /* child Jacobian row */
+
+        kn->row_map[r] = s;
+
+        /* Find C[i,j] in CSR */
+        double cij = 0.0;
+        int cij_found = 0;
+        for (int idx = C->p[i]; idx < C->p[i + 1]; idx++)
+        {
+            if (C->i[idx] == j)
+            {
+                cij = C->x[idx];
+                kn->row_scale_idx[r] = idx;
+                cij_found = 1;
+                break;
+            }
+        }
+
+        if (!cij_found || cij == 0.0)
+        {
+            kn->row_scale_idx[r] = -1; /* mark as zero row */
+            continue;
+        }
+
+        /* Copy child's column indices for row s */
+        int child_start = Jchild->p[s];
+        int child_end = Jchild->p[s + 1];
+        int child_row_nnz = child_end - child_start;
+
+        memcpy(node->jacobian->i + nnz_idx, Jchild->i + child_start,
+               (size_t) child_row_nnz * sizeof(int));
+        nnz_idx += child_row_nnz;
+    }
+    node->jacobian->p[out_size] = nnz_idx;
+    node->jacobian->nnz = nnz_idx;
+    assert(nnz_idx == total_nnz);
+}
+
+/* ------------------------------------------------------------------ */
+/*                      Jacobian evaluation                            */
+/* ------------------------------------------------------------------ */
+static void eval_jacobian(expr *node)
+{
+    kron_left_expr *kn = (kron_left_expr *) node;
+    expr *child = node->left;
+    CSR_Matrix *C = kn->C;
+    CSR_Matrix *Jchild = child->jacobian;
+    CSR_Matrix *J = node->jacobian;
+    int out_size = node->size;
+
+    /* Evaluate child's Jacobian */
+    child->eval_jacobian(child);
+
+    /* Fill values: each active row r copies child row s scaled by C[i,j] */
+    for (int r = 0; r < out_size; r++)
+    {
+        int j_start = J->p[r];
+        int j_end = J->p[r + 1];
+        int row_nnz = j_end - j_start;
+        if (row_nnz == 0) continue;
+
+        int s = kn->row_map[r];
+        int c_idx = kn->row_scale_idx[r];
+        double cij = C->x[c_idx];
+
+        int child_start = Jchild->p[s];
+
+        for (int t = 0; t < row_nnz; t++)
+        {
+            J->x[j_start + t] = cij * Jchild->x[child_start + t];
+        }
+    }
+}
+
+/* ------------------------------------------------------------------ */
+/*                  Weighted-sum Hessian initialization                 */
+/* ------------------------------------------------------------------ */
+static void wsum_hess_init(expr *node)
+{
+    expr *child = node->left;
+
+    /* Initialize child's Hessian */
+    child->wsum_hess_init(child);
+
+    /* kron_left is linear in X, so Hessian has same sparsity as child */
+    node->wsum_hess =
+        new_csr_matrix(node->n_vars, node->n_vars, child->wsum_hess->nnz);
+    memcpy(node->wsum_hess->p, child->wsum_hess->p,
+           (size_t) (node->n_vars + 1) * sizeof(int));
+    memcpy(node->wsum_hess->i, child->wsum_hess->i,
+           (size_t) child->wsum_hess->nnz * sizeof(int));
+
+    /* Allocate workspace for reverse-mode weight accumulation */
+    node->dwork = (double *) calloc((size_t) child->size, sizeof(double));
+}
+
+/* ------------------------------------------------------------------ */
+/*                  Weighted-sum Hessian evaluation                     */
+/* ------------------------------------------------------------------ */
+static void eval_wsum_hess(expr *node, const double *w)
+{
+    kron_left_expr *kn = (kron_left_expr *) node;
+    expr *child = node->left;
+    CSR_Matrix *C = kn->C;
+    int m = kn->m, p = kn->p, q = kn->q;
+    int mp = m * p;
+    int child_size = child->size;
+
+    /*
+     * Reverse mode: w_child[s] = sum over active r mapping to s of (C[i,j] * w[r])
+     * This is the adjoint of the forward pass.
+     */
+    memset(node->dwork, 0, (size_t) child_size * sizeof(double));
+
+    for (int i = 0; i < m; i++)
+    {
+        for (int idx = C->p[i]; idx < C->p[i + 1]; idx++)
+        {
+            int j = C->i[idx];
+            double cij = C->x[idx];
+
+            for (int l = 0; l < q; l++)
+            {
+                for (int k = 0; k < p; k++)
+                {
+                    int r = (j * q + l) * mp + i * p + k;
+                    int s = l * p + k;
+                    node->dwork[s] += cij * w[r];
+                }
+            }
+        }
+    }
+
+    /* Delegate to child */
+    child->eval_wsum_hess(child, node->dwork);
+    memcpy(node->wsum_hess->x, child->wsum_hess->x,
+           (size_t) node->wsum_hess->nnz * sizeof(double));
+}
+
+/* ------------------------------------------------------------------ */
+/*                          Cleanup                                    */
+/* ------------------------------------------------------------------ */
+static void free_type_data(expr *node)
+{
+    kron_left_expr *kn = (kron_left_expr *) node;
+    free_csr_matrix(kn->C);
+    free(kn->row_map);
+    free(kn->row_scale_idx);
+    kn->C = NULL;
+    kn->row_map = NULL;
+    kn->row_scale_idx = NULL;
+}
+
+/* ------------------------------------------------------------------ */
+/*                        Constructor                                  */
+/* ------------------------------------------------------------------ */
+expr *new_kron_left(expr *child, const CSR_Matrix *C, int p, int q)
+{
+    int m = C->m;
+    int n = C->n;
+
+    /* Verify child dimensions */
+    if (child->size != p * q)
+    {
+        fprintf(stderr,
+                "Error in new_kron_left: child size %d != p*q = %d*%d = %d\n",
+                child->size, p, q, p * q);
+        exit(1);
+    }
+
+    /* Output: kron(C, X) has shape (m*p, n*q) */
+    int d1 = m * p;
+    int d2 = n * q;
+
+    kron_left_expr *kn = (kron_left_expr *) calloc(1, sizeof(kron_left_expr));
+    expr *node = &kn->base;
+    init_expr(node, d1, d2, child->n_vars, forward, jacobian_init, eval_jacobian,
+              is_affine, wsum_hess_init, eval_wsum_hess, free_type_data);
+
+    node->left = child;
+    expr_retain(child);
+
+    kn->m = m;
+    kn->n = n;
+    kn->p = p;
+    kn->q = q;
+    kn->C = new_csr(C);
+    kn->row_map = NULL;
+    kn->row_scale_idx = NULL;
+
+    return node;
+}