From 7a09d56f3e98f8a78a8a098728b61107621d00d2 Mon Sep 17 00:00:00 2001
From: William Zijie Zhang <william@gridmatic.com>
Date: Mon, 5 Jan 2026 00:03:19 -0500
Subject: [PATCH 01/27] adds problem struct

---
 python/bindings.c             | 257 ++++++++++++++++++++++++
 python/convert.py             |  47 ++++-
 python/tests/convert_tests.py | 358 ----------------------------------
 tests/all_tests.c             |   7 +
 4 files changed, 309 insertions(+), 360 deletions(-)
 delete mode 100644 python/tests/convert_tests.py

diff --git a/python/bindings.c b/python/bindings.c
index 68e4b76..dce1477 100644
--- a/python/bindings.c
+++ b/python/bindings.c
@@ -6,9 +6,11 @@
 #include "affine.h"
 #include "elementwise_univariate.h"
 #include "expr.h"
+#include "problem.h"
 
 // Capsule name for expr* pointers
 #define EXPR_CAPSULE_NAME "DNLP_EXPR"
+#define PROBLEM_CAPSULE_NAME "DNLP_PROBLEM"
 
 static int numpy_initialized = 0;
 
@@ -237,6 +239,256 @@ static PyObject *py_jacobian(PyObject *self, PyObject *args)
     return Py_BuildValue("(OOO(ii))", data, indices, indptr, jac->m, jac->n);
 }
 
+/* ========== Problem bindings ========== */
+
+static void problem_capsule_destructor(PyObject *capsule)
+{
+    problem *prob = (problem *) PyCapsule_GetPointer(capsule, PROBLEM_CAPSULE_NAME);
+    if (prob)
+    {
+        free_problem(prob);
+    }
+}
+
+static PyObject *py_make_problem(PyObject *self, PyObject *args)
+{
+    PyObject *obj_capsule;
+    PyObject *constraints_list;
+    if (!PyArg_ParseTuple(args, "OO", &obj_capsule, &constraints_list))
+    {
+        return NULL;
+    }
+
+    expr *objective = (expr *) PyCapsule_GetPointer(obj_capsule, EXPR_CAPSULE_NAME);
+    if (!objective)
+    {
+        PyErr_SetString(PyExc_ValueError, "invalid objective capsule");
+        return NULL;
+    }
+
+    if (!PyList_Check(constraints_list))
+    {
+        PyErr_SetString(PyExc_TypeError, "constraints must be a list");
+        return NULL;
+    }
+
+    Py_ssize_t n_constraints = PyList_Size(constraints_list);
+    expr **constraints = NULL;
+    if (n_constraints > 0)
+    {
+        constraints = (expr **) malloc(n_constraints * sizeof(expr *));
+        if (!constraints)
+        {
+            PyErr_NoMemory();
+            return NULL;
+        }
+        for (Py_ssize_t i = 0; i < n_constraints; i++)
+        {
+            PyObject *c_capsule = PyList_GetItem(constraints_list, i);
+            constraints[i] =
+                (expr *) PyCapsule_GetPointer(c_capsule, EXPR_CAPSULE_NAME);
+            if (!constraints[i])
+            {
+                free(constraints);
+                PyErr_SetString(PyExc_ValueError, "invalid constraint capsule");
+                return NULL;
+            }
+        }
+    }
+
+    problem *prob = new_problem(objective, constraints, (int) n_constraints);
+    free(constraints);
+
+    if (!prob)
+    {
+        PyErr_SetString(PyExc_RuntimeError, "failed to create problem");
+        return NULL;
+    }
+
+    return PyCapsule_New(prob, PROBLEM_CAPSULE_NAME, problem_capsule_destructor);
+}
+
+static PyObject *py_problem_allocate(PyObject *self, PyObject *args)
+{
+    PyObject *prob_capsule;
+    PyObject *u_obj;
+    if (!PyArg_ParseTuple(args, "OO", &prob_capsule, &u_obj))
+    {
+        return NULL;
+    }
+
+    problem *prob =
+        (problem *) PyCapsule_GetPointer(prob_capsule, PROBLEM_CAPSULE_NAME);
+    if (!prob)
+    {
+        PyErr_SetString(PyExc_ValueError, "invalid problem capsule");
+        return NULL;
+    }
+
+    PyArrayObject *u_array =
+        (PyArrayObject *) PyArray_FROM_OTF(u_obj, NPY_DOUBLE, NPY_ARRAY_IN_ARRAY);
+    if (!u_array)
+    {
+        return NULL;
+    }
+
+    problem_allocate(prob, (const double *) PyArray_DATA(u_array));
+
+    Py_DECREF(u_array);
+    Py_RETURN_NONE;
+}
+
+static PyObject *py_problem_forward(PyObject *self, PyObject *args)
+{
+    PyObject *prob_capsule;
+    PyObject *u_obj;
+    if (!PyArg_ParseTuple(args, "OO", &prob_capsule, &u_obj))
+    {
+        return NULL;
+    }
+
+    problem *prob =
+        (problem *) PyCapsule_GetPointer(prob_capsule, PROBLEM_CAPSULE_NAME);
+    if (!prob)
+    {
+        PyErr_SetString(PyExc_ValueError, "invalid problem capsule");
+        return NULL;
+    }
+
+    PyArrayObject *u_array =
+        (PyArrayObject *) PyArray_FROM_OTF(u_obj, NPY_DOUBLE, NPY_ARRAY_IN_ARRAY);
+    if (!u_array)
+    {
+        return NULL;
+    }
+
+    double obj_val = problem_forward(prob, (const double *) PyArray_DATA(u_array));
+
+    // Create constraint values array
+    PyObject *constraint_vals = NULL;
+    if (prob->total_constraint_size > 0)
+    {
+        npy_intp size = prob->total_constraint_size;
+        constraint_vals = PyArray_SimpleNew(1, &size, NPY_DOUBLE);
+        if (!constraint_vals)
+        {
+            Py_DECREF(u_array);
+            return NULL;
+        }
+        memcpy(PyArray_DATA((PyArrayObject *) constraint_vals), prob->constraint_values,
+               size * sizeof(double));
+    }
+    else
+    {
+        npy_intp size = 0;
+        constraint_vals = PyArray_SimpleNew(1, &size, NPY_DOUBLE);
+    }
+
+    Py_DECREF(u_array);
+    return Py_BuildValue("(dO)", obj_val, constraint_vals);
+}
+
+static PyObject *py_problem_gradient(PyObject *self, PyObject *args)
+{
+    PyObject *prob_capsule;
+    PyObject *u_obj;
+    if (!PyArg_ParseTuple(args, "OO", &prob_capsule, &u_obj))
+    {
+        return NULL;
+    }
+
+    problem *prob =
+        (problem *) PyCapsule_GetPointer(prob_capsule, PROBLEM_CAPSULE_NAME);
+    if (!prob)
+    {
+        PyErr_SetString(PyExc_ValueError, "invalid problem capsule");
+        return NULL;
+    }
+
+    PyArrayObject *u_array =
+        (PyArrayObject *) PyArray_FROM_OTF(u_obj, NPY_DOUBLE, NPY_ARRAY_IN_ARRAY);
+    if (!u_array)
+    {
+        return NULL;
+    }
+
+    double *grad = problem_gradient(prob, (const double *) PyArray_DATA(u_array));
+
+    npy_intp size = prob->n_vars;
+    PyObject *out = PyArray_SimpleNew(1, &size, NPY_DOUBLE);
+    if (!out)
+    {
+        Py_DECREF(u_array);
+        return NULL;
+    }
+    memcpy(PyArray_DATA((PyArrayObject *) out), grad, size * sizeof(double));
+
+    Py_DECREF(u_array);
+    return out;
+}
+
+static PyObject *py_problem_jacobian(PyObject *self, PyObject *args)
+{
+    PyObject *prob_capsule;
+    PyObject *u_obj;
+    if (!PyArg_ParseTuple(args, "OO", &prob_capsule, &u_obj))
+    {
+        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->n_constraints == 0)
+    {
+        // Return empty CSR components
+        npy_intp zero = 0;
+        npy_intp one = 1;
+        PyObject *data = PyArray_SimpleNew(1, &zero, NPY_DOUBLE);
+        PyObject *indices = PyArray_SimpleNew(1, &zero, NPY_INT32);
+        PyObject *indptr = PyArray_SimpleNew(1, &one, NPY_INT32);
+        ((int *) PyArray_DATA((PyArrayObject *) indptr))[0] = 0;
+        return Py_BuildValue("(OOO(ii))", data, indices, indptr, 0, prob->n_vars);
+    }
+
+    PyArrayObject *u_array =
+        (PyArrayObject *) PyArray_FROM_OTF(u_obj, NPY_DOUBLE, NPY_ARRAY_IN_ARRAY);
+    if (!u_array)
+    {
+        return NULL;
+    }
+
+    CSR_Matrix *jac = problem_jacobian(prob, (const double *) PyArray_DATA(u_array));
+
+    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);
+        Py_DECREF(u_array);
+        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));
+
+    Py_DECREF(u_array);
+    return Py_BuildValue("(OOO(ii))", data, indices, indptr, jac->m, jac->n);
+}
+
 static PyMethodDef DNLPMethods[] = {
     {"make_variable", py_make_variable, METH_VARARGS, "Create variable node"},
     {"make_log", py_make_log, METH_VARARGS, "Create log node"},
@@ -245,6 +497,11 @@ static PyMethodDef DNLPMethods[] = {
     {"make_sum", py_make_sum, METH_VARARGS, "Create sum node"},
     {"forward", py_forward, METH_VARARGS, "Run forward pass and return values"},
     {"jacobian", py_jacobian, METH_VARARGS, "Compute jacobian and return CSR components"},
+    {"make_problem", py_make_problem, METH_VARARGS, "Create problem from objective and constraints"},
+    {"problem_allocate", py_problem_allocate, METH_VARARGS, "Allocate problem resources"},
+    {"problem_forward", py_problem_forward, METH_VARARGS, "Evaluate objective and constraints"},
+    {"problem_gradient", py_problem_gradient, METH_VARARGS, "Compute objective gradient"},
+    {"problem_jacobian", py_problem_jacobian, METH_VARARGS, "Compute constraint jacobian"},
     {NULL, NULL, 0, NULL}};
 
 static struct PyModuleDef dnlp_module = {PyModuleDef_HEAD_INIT, "DNLP_diff_engine",
diff --git a/python/convert.py b/python/convert.py
index 7a1be62..bf8ccd4 100644
--- a/python/convert.py
+++ b/python/convert.py
@@ -1,4 +1,5 @@
 import cvxpy as cp
+import numpy as np
 from scipy import sparse
 from cvxpy.reductions.inverse_data import InverseData
 import DNLP_diff_engine as diffengine
@@ -80,9 +81,9 @@ def _convert_expr(expr, var_dict: dict):
     raise NotImplementedError(f"Atom '{atom_name}' not supported")
 
 
-def convert_problem(problem: cp.Problem) -> tuple:
+def convert_expressions(problem: cp.Problem) -> tuple:
     """
-    Convert CVXPY Problem to C expressions.
+    Convert CVXPY Problem to C expressions (low-level).
 
     Args:
         problem: CVXPY Problem object
@@ -103,3 +104,45 @@ def convert_problem(problem: cp.Problem) -> tuple:
         c_constraints.append(c_expr)
 
     return c_objective, c_constraints
+
+
+def convert_problem(problem: cp.Problem) -> "Problem":
+    """
+    Convert CVXPY Problem to C problem struct.
+
+    Args:
+        problem: CVXPY Problem object
+
+    Returns:
+        Problem wrapper around C problem struct
+    """
+    return Problem(problem)
+
+
+class Problem:
+    """Wrapper around C problem struct for CVXPY problems."""
+
+    def __init__(self, cvxpy_problem: cp.Problem):
+        var_dict = build_variable_dict(cvxpy_problem.variables())
+        c_obj = _convert_expr(cvxpy_problem.objective.expr, var_dict)
+        c_constraints = [_convert_expr(c.expr, var_dict) for c in cvxpy_problem.constraints]
+        self._capsule = diffengine.make_problem(c_obj, c_constraints)
+        self._allocated = False
+
+    def allocate(self, u: np.ndarray):
+        """Allocate internal buffers. Must be called before forward/gradient/jacobian."""
+        diffengine.problem_allocate(self._capsule, u)
+        self._allocated = True
+
+    def forward(self, u: np.ndarray) -> tuple[float, np.ndarray]:
+        """Evaluate objective and constraints. Returns (obj_value, constraint_values)."""
+        return diffengine.problem_forward(self._capsule, u)
+
+    def gradient(self, u: np.ndarray) -> np.ndarray:
+        """Compute gradient of objective. Returns gradient array."""
+        return diffengine.problem_gradient(self._capsule, u)
+
+    def jacobian(self, u: np.ndarray) -> sparse.csr_matrix:
+        """Compute jacobian of constraints. Returns scipy CSR matrix."""
+        data, indices, indptr, shape = diffengine.problem_jacobian(self._capsule, u)
+        return sparse.csr_matrix((data, indices, indptr), shape=shape)
diff --git a/python/tests/convert_tests.py b/python/tests/convert_tests.py
deleted file mode 100644
index 0bed90a..0000000
--- a/python/tests/convert_tests.py
+++ /dev/null
@@ -1,358 +0,0 @@
-import cvxpy as cp
-import numpy as np
-import DNLP_diff_engine as diffengine
-from convert import convert_problem, get_jacobian
-
-
-def test_sum_log():
-    """Test sum(log(x)) forward and jacobian."""
-    x = cp.Variable(4)
-    problem = cp.Problem(cp.Minimize(cp.sum(cp.log(x))))
-    c_obj, _ = convert_problem(problem)
-
-    test_values = np.array([1.0, 2.0, 3.0, 4.0])
-
-    # Forward
-    result = diffengine.forward(c_obj, test_values)
-    expected = np.sum(np.log(test_values))
-    assert np.allclose(result, expected)
-
-    # Jacobian: d/dx sum(log(x)) = [1/x_1, 1/x_2, ...]
-    jac = get_jacobian(c_obj, test_values)
-    expected_jac = (1.0 / test_values).reshape(1, -1)
-    assert np.allclose(jac.toarray(), expected_jac)
-
-
-def test_sum_exp():
-    """Test sum(exp(x)) forward and jacobian."""
-    x = cp.Variable(3)
-    problem = cp.Problem(cp.Minimize(cp.sum(cp.exp(x))))
-    c_obj, _ = convert_problem(problem)
-
-    test_values = np.array([0.0, 1.0, 2.0])
-
-    # Forward
-    result = diffengine.forward(c_obj, test_values)
-    expected = np.sum(np.exp(test_values))
-    assert np.allclose(result, expected)
-
-    # Jacobian: d/dx sum(exp(x)) = [exp(x_1), exp(x_2), ...]
-    jac = get_jacobian(c_obj, test_values)
-    expected_jac = np.exp(test_values).reshape(1, -1)
-    assert np.allclose(jac.toarray(), expected_jac)
-
-
-def test_two_variables_elementwise_add():
-    """Test sum(log(x + y)) - elementwise after add."""
-    x = cp.Variable(2)
-    y = cp.Variable(2)
-    problem = cp.Problem(cp.Minimize(cp.sum(cp.log(x + y))))
-    c_obj, _ = convert_problem(problem)
-
-    test_values = np.array([1.0, 2.0, 3.0, 4.0])
-
-    # Forward
-    result = diffengine.forward(c_obj, test_values)
-    expected = np.sum(np.log(np.array([1+3, 2+4])))
-    assert np.allclose(result, expected)
-
-    # TODO: Jacobian for elementwise(add(...)) patterns not yet supported
-
-
-def test_variable_reuse():
-    """Test sum(log(x) + exp(x)) - same variable used twice."""
-    x = cp.Variable(2)
-    problem = cp.Problem(cp.Minimize(cp.sum(cp.log(x) + cp.exp(x))))
-    c_obj, _ = convert_problem(problem)
-
-    test_values = np.array([1.0, 2.0])
-
-    # Forward
-    result = diffengine.forward(c_obj, test_values)
-    expected = np.sum(np.log(test_values) + np.exp(test_values))
-    assert np.allclose(result, expected)
-
-    # Jacobian: d/dx_i = 1/x_i + exp(x_i)
-    jac = get_jacobian(c_obj, test_values)
-    expected_jac = (1.0 / test_values + np.exp(test_values)).reshape(1, -1)
-    assert np.allclose(jac.toarray(), expected_jac)
-
-
-def test_four_variables_elementwise_add():
-    """Test sum(log(a + b) + exp(c + d)) - elementwise after add."""
-    a = cp.Variable(3)
-    b = cp.Variable(3)
-    c = cp.Variable(3)
-    d = cp.Variable(3)
-    problem = cp.Problem(cp.Minimize(cp.sum(cp.log(a + b) + cp.exp(c + d))))
-    c_obj, _ = convert_problem(problem)
-
-    a_vals = np.array([1.0, 2.0, 3.0])
-    b_vals = np.array([0.5, 1.0, 1.5])
-    c_vals = np.array([0.1, 0.2, 0.3])
-    d_vals = np.array([0.1, 0.1, 0.1])
-    test_values = np.concatenate([a_vals, b_vals, c_vals, d_vals])
-
-    # Forward
-    result = diffengine.forward(c_obj, test_values)
-    expected = np.sum(np.log(a_vals + b_vals) + np.exp(c_vals + d_vals))
-    assert np.allclose(result, expected)
-
-    # TODO: Jacobian for elementwise(add(...)) patterns not yet supported
-
-
-def test_deep_nesting():
-    """Test sum(log(exp(log(exp(x))))) - deeply nested elementwise."""
-    x = cp.Variable(4)
-    problem = cp.Problem(cp.Minimize(cp.sum(cp.log(cp.exp(cp.log(cp.exp(x)))))))
-    c_obj, _ = convert_problem(problem)
-
-    test_values = np.array([0.5, 1.0, 1.5, 2.0])
-
-    # Forward
-    result = diffengine.forward(c_obj, test_values)
-    expected = np.sum(np.log(np.exp(np.log(np.exp(test_values)))))
-    assert np.allclose(result, expected)
-
-    # TODO: Jacobian for nested elementwise compositions not yet supported
-
-
-def test_chained_additions():
-    """Test sum(x + y + z + w) - chained additions."""
-    x = cp.Variable(2)
-    y = cp.Variable(2)
-    z = cp.Variable(2)
-    w = cp.Variable(2)
-    problem = cp.Problem(cp.Minimize(cp.sum(x + y + z + w)))
-    c_obj, _ = convert_problem(problem)
-
-    x_vals = np.array([1.0, 2.0])
-    y_vals = np.array([3.0, 4.0])
-    z_vals = np.array([5.0, 6.0])
-    w_vals = np.array([7.0, 8.0])
-    test_values = np.concatenate([x_vals, y_vals, z_vals, w_vals])
-
-    # Forward
-    result = diffengine.forward(c_obj, test_values)
-    expected = np.sum(x_vals + y_vals + z_vals + w_vals)
-    assert np.allclose(result, expected)
-
-    # TODO: Jacobian for sum(add(...)) patterns not yet supported
-
-
-def test_variable_used_multiple_times():
-    """Test sum(log(x) + exp(x) + x) - variable used 3+ times."""
-    x = cp.Variable(3)
-    problem = cp.Problem(cp.Minimize(cp.sum(cp.log(x) + cp.exp(x) + x)))
-    c_obj, _ = convert_problem(problem)
-
-    test_values = np.array([1.0, 2.0, 3.0])
-
-    # Forward
-    result = diffengine.forward(c_obj, test_values)
-    expected = np.sum(np.log(test_values) + np.exp(test_values) + test_values)
-    assert np.allclose(result, expected)
-
-    # TODO: Jacobian for expressions with sum(variable) not yet supported
-
-
-def test_larger_variable():
-    """Test sum(log(x)) with larger variable (100 elements)."""
-    x = cp.Variable(100)
-    problem = cp.Problem(cp.Minimize(cp.sum(cp.log(x))))
-    c_obj, _ = convert_problem(problem)
-
-    test_values = np.linspace(1.0, 10.0, 100)
-
-    # Forward
-    result = diffengine.forward(c_obj, test_values)
-    expected = np.sum(np.log(test_values))
-    assert np.allclose(result, expected)
-
-    # Jacobian
-    jac = get_jacobian(c_obj, test_values)
-    expected_jac = (1.0 / test_values).reshape(1, -1)
-    assert np.allclose(jac.toarray(), expected_jac)
-
-
-def test_matrix_variable():
-    """Test sum(log(X)) with 2D matrix variable (3x4)."""
-    X = cp.Variable((3, 4))
-    problem = cp.Problem(cp.Minimize(cp.sum(cp.log(X))))
-    c_obj, _ = convert_problem(problem)
-
-    test_values = np.arange(1.0, 13.0)  # 12 elements
-
-    # Forward
-    result = diffengine.forward(c_obj, test_values)
-    expected = np.sum(np.log(test_values))
-    assert np.allclose(result, expected)
-
-    # Jacobian
-    jac = get_jacobian(c_obj, test_values)
-    expected_jac = (1.0 / test_values).reshape(1, -1)
-    assert np.allclose(jac.toarray(), expected_jac)
-
-
-def test_mixed_sizes():
-    """Test sum(log(a)) + sum(log(b)) + sum(log(c)) with different sized variables."""
-    a = cp.Variable(2)
-    b = cp.Variable(5)
-    c = cp.Variable(3)
-    problem = cp.Problem(cp.Minimize(cp.sum(cp.log(a)) + cp.sum(cp.log(b)) + cp.sum(cp.log(c))))
-    c_obj, _ = convert_problem(problem)
-
-    a_vals = np.array([1.0, 2.0])
-    b_vals = np.array([1.0, 2.0, 3.0, 4.0, 5.0])
-    c_vals = np.array([1.0, 2.0, 3.0])
-    test_values = np.concatenate([a_vals, b_vals, c_vals])
-
-    # Forward
-    result = diffengine.forward(c_obj, test_values)
-    expected = np.sum(np.log(a_vals)) + np.sum(np.log(b_vals)) + np.sum(np.log(c_vals))
-    assert np.allclose(result, expected)
-
-    # Jacobian
-    jac = get_jacobian(c_obj, test_values)
-    expected_jac = (1.0 / test_values).reshape(1, -1)
-    assert np.allclose(jac.toarray(), expected_jac)
-
-
-def test_multiple_variables_log_exp():
-    """Test sum(log(a)) + sum(log(b)) + sum(exp(c)) + sum(exp(d))."""
-    a = cp.Variable(2)
-    b = cp.Variable(2)
-    c = cp.Variable(2)
-    d = cp.Variable(2)
-    obj = cp.sum(cp.log(a)) + cp.sum(cp.log(b)) + cp.sum(cp.exp(c)) + cp.sum(cp.exp(d))
-    problem = cp.Problem(cp.Minimize(obj))
-    c_obj, _ = convert_problem(problem)
-
-    a_vals = np.array([1.0, 2.0])
-    b_vals = np.array([0.5, 1.0])
-    c_vals = np.array([0.1, 0.2])
-    d_vals = np.array([0.1, 0.1])
-    test_values = np.concatenate([a_vals, b_vals, c_vals, d_vals])
-
-    # Forward
-    result = diffengine.forward(c_obj, test_values)
-    expected = (np.sum(np.log(a_vals)) + np.sum(np.log(b_vals)) +
-                np.sum(np.exp(c_vals)) + np.sum(np.exp(d_vals)))
-    assert np.allclose(result, expected)
-
-    # Jacobian
-    jac = get_jacobian(c_obj, test_values)
-    df_da = 1.0 / a_vals
-    df_db = 1.0 / b_vals
-    df_dc = np.exp(c_vals)
-    df_dd = np.exp(d_vals)
-    expected_jac = np.concatenate([df_da, df_db, df_dc, df_dd]).reshape(1, -1)
-    assert np.allclose(jac.toarray(), expected_jac)
-
-
-def test_two_variables_separate_sums():
-    """Test sum(log(x) + log(y)) - two variables with separate elementwise ops."""
-    x = cp.Variable(2)
-    y = cp.Variable(2)
-    problem = cp.Problem(cp.Minimize(cp.sum(cp.log(x) + cp.log(y))))
-    c_obj, _ = convert_problem(problem)
-
-    test_values = np.array([1.0, 2.0, 3.0, 4.0])
-
-    # Forward
-    result = diffengine.forward(c_obj, test_values)
-    expected = np.sum(np.log(test_values[:2]) + np.log(test_values[2:]))
-    assert np.allclose(result, expected)
-
-    # Jacobian
-    jac = get_jacobian(c_obj, test_values)
-    expected_jac = np.array([[1/1, 1/2, 1/3, 1/4]])
-    assert np.allclose(jac.toarray(), expected_jac)
-
-
-def test_complex_objective_elementwise_add():
-    """Test sum(log(x + y)) + sum(exp(y + z)) + sum(log(z + x)) - elementwise after add."""
-    x = cp.Variable(3)
-    y = cp.Variable(3)
-    z = cp.Variable(3)
-    obj = cp.sum(cp.log(x + y)) + cp.sum(cp.exp(y + z)) + cp.sum(cp.log(z + x))
-    problem = cp.Problem(cp.Minimize(obj))
-    c_obj, _ = convert_problem(problem)
-
-    x_vals = np.array([1.0, 2.0, 3.0])
-    y_vals = np.array([0.5, 1.0, 1.5])
-    z_vals = np.array([0.2, 0.3, 0.4])
-    test_values = np.concatenate([x_vals, y_vals, z_vals])
-
-    # Forward
-    result = diffengine.forward(c_obj, test_values)
-    expected = (np.sum(np.log(x_vals + y_vals)) +
-                np.sum(np.exp(y_vals + z_vals)) +
-                np.sum(np.log(z_vals + x_vals)))
-    assert np.allclose(result, expected)
-
-    # TODO: Jacobian for elementwise(add(...)) patterns not yet supported
-
-
-def test_complex_objective_no_add():
-    """Test sum(log(x) + exp(y) + log(z)) - multiple elementwise ops without add composition."""
-    x = cp.Variable(2)
-    y = cp.Variable(2)
-    z = cp.Variable(2)
-    obj = cp.sum(cp.log(x) + cp.exp(y) + cp.log(z))
-    problem = cp.Problem(cp.Minimize(obj))
-    c_obj, _ = convert_problem(problem)
-
-    x_vals = np.array([1.0, 2.0])
-    y_vals = np.array([0.5, 1.0])
-    z_vals = np.array([0.2, 0.3])
-    test_values = np.concatenate([x_vals, y_vals, z_vals])
-
-    # Forward
-    result = diffengine.forward(c_obj, test_values)
-    expected = np.sum(np.log(x_vals) + np.exp(y_vals) + np.log(z_vals))
-    assert np.allclose(result, expected)
-
-    # Jacobian
-    jac = get_jacobian(c_obj, test_values)
-    df_dx = 1.0 / x_vals
-    df_dy = np.exp(y_vals)
-    df_dz = 1.0 / z_vals
-    expected_jac = np.concatenate([df_dx, df_dy, df_dz]).reshape(1, -1)
-    assert np.allclose(jac.toarray(), expected_jac)
-
-
-def test_log_exp_identity():
-    """Test sum(log(exp(x))) = sum(x) identity - nested elementwise."""
-    x = cp.Variable(5)
-    problem = cp.Problem(cp.Minimize(cp.sum(cp.log(cp.exp(x)))))
-    c_obj, _ = convert_problem(problem)
-
-    test_values = np.array([-1.0, 0.0, 1.0, 2.0, 3.0])
-
-    # Forward
-    result = diffengine.forward(c_obj, test_values)
-    expected = np.sum(test_values)  # log(exp(x)) = x
-    assert np.allclose(result, expected)
-
-    # TODO: Jacobian for nested elementwise compositions not yet supported
-
-
-if __name__ == "__main__":
-    test_sum_log()
-    test_sum_exp()
-    test_two_variables_elementwise_add()
-    test_variable_reuse()
-    test_four_variables_elementwise_add()
-    test_deep_nesting()
-    test_chained_additions()
-    test_variable_used_multiple_times()
-    test_larger_variable()
-    test_matrix_variable()
-    test_mixed_sizes()
-    test_multiple_variables_log_exp()
-    test_two_variables_separate_sums()
-    test_complex_objective_elementwise_add()
-    test_complex_objective_no_add()
-    test_log_exp_identity()
-    print("All tests passed!")
diff --git a/tests/all_tests.c b/tests/all_tests.c
index d57c183..a7cc644 100644
--- a/tests/all_tests.c
+++ b/tests/all_tests.c
@@ -32,6 +32,7 @@
 #include "wsum_hess/test_hstack.h"
 #include "wsum_hess/test_rel_entr.h"
 #include "wsum_hess/test_sum.h"
+#include "problem/test_problem.h"
 
 int main(void)
 {
@@ -118,6 +119,12 @@ int main(void)
     mu_run_test(test_ATA_alloc_random, tests_run);
     mu_run_test(test_ATA_alloc_random2, tests_run);
 
+    printf("\n--- Problem Struct Tests ---\n");
+    mu_run_test(test_problem_new_free, tests_run);
+    mu_run_test(test_problem_forward, tests_run);
+    mu_run_test(test_problem_gradient, tests_run);
+    mu_run_test(test_problem_jacobian, tests_run);
+
     printf("\n=== All %d tests passed ===\n", tests_run);
 
     return 0;

From acdd79e8cab3268e3142aa9868d7cbca2aeff0a8 Mon Sep 17 00:00:00 2001
From: William Zijie Zhang <william@gridmatic.com>
Date: Mon, 5 Jan 2026 00:03:55 -0500
Subject: [PATCH 02/27] adds newly added files

---
 include/problem.h                    |  29 +++
 python/tests/test_basic_convert.py   | 359 +++++++++++++++++++++++++++
 python/tests/test_problem_convert.py | 337 +++++++++++++++++++++++++
 src/problem.c                        | 179 +++++++++++++
 tests/problem/test_problem.h         | 162 ++++++++++++
 5 files changed, 1066 insertions(+)
 create mode 100644 include/problem.h
 create mode 100644 python/tests/test_basic_convert.py
 create mode 100644 python/tests/test_problem_convert.py
 create mode 100644 src/problem.c
 create mode 100644 tests/problem/test_problem.h

diff --git a/include/problem.h b/include/problem.h
new file mode 100644
index 0000000..97b0e9f
--- /dev/null
+++ b/include/problem.h
@@ -0,0 +1,29 @@
+#ifndef PROBLEM_H
+#define PROBLEM_H
+
+#include "expr.h"
+#include "utils/CSR_Matrix.h"
+
+typedef struct problem
+{
+    expr *objective;
+    expr **constraints;
+    int n_constraints;
+    int n_vars;
+    int total_constraint_size;
+
+    /* Allocated by problem_allocate */
+    double *constraint_values;
+    double *gradient_values;
+    CSR_Matrix *stacked_jac;
+} problem;
+
+problem *new_problem(expr *objective, expr **constraints, int n_constraints);
+void problem_allocate(problem *prob, const double *u);
+void free_problem(problem *prob);
+
+double problem_forward(problem *prob, const double *u);
+double *problem_gradient(problem *prob, const double *u);
+CSR_Matrix *problem_jacobian(problem *prob, const double *u);
+
+#endif
diff --git a/python/tests/test_basic_convert.py b/python/tests/test_basic_convert.py
new file mode 100644
index 0000000..40318d8
--- /dev/null
+++ b/python/tests/test_basic_convert.py
@@ -0,0 +1,359 @@
+import cvxpy as cp
+import numpy as np
+from scipy import sparse
+import DNLP_diff_engine as diffengine
+from convert import convert_expressions, get_jacobian
+
+
+def test_sum_log():
+    """Test sum(log(x)) forward and jacobian."""
+    x = cp.Variable(4)
+    problem = cp.Problem(cp.Minimize(cp.sum(cp.log(x))))
+    c_obj, _ = convert_expressions(problem)
+
+    test_values = np.array([1.0, 2.0, 3.0, 4.0])
+
+    # Forward
+    result = diffengine.forward(c_obj, test_values)
+    expected = np.sum(np.log(test_values))
+    assert np.allclose(result, expected)
+
+    # Jacobian: d/dx sum(log(x)) = [1/x_1, 1/x_2, ...]
+    jac = get_jacobian(c_obj, test_values)
+    expected_jac = (1.0 / test_values).reshape(1, -1)
+    assert np.allclose(jac.toarray(), expected_jac)
+
+
+def test_sum_exp():
+    """Test sum(exp(x)) forward and jacobian."""
+    x = cp.Variable(3)
+    problem = cp.Problem(cp.Minimize(cp.sum(cp.exp(x))))
+    c_obj, _ = convert_expressions(problem)
+
+    test_values = np.array([0.0, 1.0, 2.0])
+
+    # Forward
+    result = diffengine.forward(c_obj, test_values)
+    expected = np.sum(np.exp(test_values))
+    assert np.allclose(result, expected)
+
+    # Jacobian: d/dx sum(exp(x)) = [exp(x_1), exp(x_2), ...]
+    jac = get_jacobian(c_obj, test_values)
+    expected_jac = np.exp(test_values).reshape(1, -1)
+    assert np.allclose(jac.toarray(), expected_jac)
+
+
+def test_two_variables_elementwise_add():
+    """Test sum(log(x + y)) - elementwise after add."""
+    x = cp.Variable(2)
+    y = cp.Variable(2)
+    problem = cp.Problem(cp.Minimize(cp.sum(cp.log(x + y))))
+    c_obj, _ = convert_expressions(problem)
+
+    test_values = np.array([1.0, 2.0, 3.0, 4.0])
+
+    # Forward
+    result = diffengine.forward(c_obj, test_values)
+    expected = np.sum(np.log(np.array([1+3, 2+4])))
+    assert np.allclose(result, expected)
+
+    # TODO: Jacobian for elementwise(add(...)) patterns not yet supported
+
+
+def test_variable_reuse():
+    """Test sum(log(x) + exp(x)) - same variable used twice."""
+    x = cp.Variable(2)
+    problem = cp.Problem(cp.Minimize(cp.sum(cp.log(x) + cp.exp(x))))
+    c_obj, _ = convert_expressions(problem)
+
+    test_values = np.array([1.0, 2.0])
+
+    # Forward
+    result = diffengine.forward(c_obj, test_values)
+    expected = np.sum(np.log(test_values) + np.exp(test_values))
+    assert np.allclose(result, expected)
+
+    # Jacobian: d/dx_i = 1/x_i + exp(x_i)
+    jac = get_jacobian(c_obj, test_values)
+    expected_jac = (1.0 / test_values + np.exp(test_values)).reshape(1, -1)
+    assert np.allclose(jac.toarray(), expected_jac)
+
+
+def test_four_variables_elementwise_add():
+    """Test sum(log(a + b) + exp(c + d)) - elementwise after add."""
+    a = cp.Variable(3)
+    b = cp.Variable(3)
+    c = cp.Variable(3)
+    d = cp.Variable(3)
+    problem = cp.Problem(cp.Minimize(cp.sum(cp.log(a + b) + cp.exp(c + d))))
+    c_obj, _ = convert_expressions(problem)
+
+    a_vals = np.array([1.0, 2.0, 3.0])
+    b_vals = np.array([0.5, 1.0, 1.5])
+    c_vals = np.array([0.1, 0.2, 0.3])
+    d_vals = np.array([0.1, 0.1, 0.1])
+    test_values = np.concatenate([a_vals, b_vals, c_vals, d_vals])
+
+    # Forward
+    result = diffengine.forward(c_obj, test_values)
+    expected = np.sum(np.log(a_vals + b_vals) + np.exp(c_vals + d_vals))
+    assert np.allclose(result, expected)
+
+    # TODO: Jacobian for elementwise(add(...)) patterns not yet supported
+
+
+def test_deep_nesting():
+    """Test sum(log(exp(log(exp(x))))) - deeply nested elementwise."""
+    x = cp.Variable(4)
+    problem = cp.Problem(cp.Minimize(cp.sum(cp.log(cp.exp(cp.log(cp.exp(x)))))))
+    c_obj, _ = convert_expressions(problem)
+
+    test_values = np.array([0.5, 1.0, 1.5, 2.0])
+
+    # Forward
+    result = diffengine.forward(c_obj, test_values)
+    expected = np.sum(np.log(np.exp(np.log(np.exp(test_values)))))
+    assert np.allclose(result, expected)
+
+    # TODO: Jacobian for nested elementwise compositions not yet supported
+
+
+def test_chained_additions():
+    """Test sum(x + y + z + w) - chained additions."""
+    x = cp.Variable(2)
+    y = cp.Variable(2)
+    z = cp.Variable(2)
+    w = cp.Variable(2)
+    problem = cp.Problem(cp.Minimize(cp.sum(x + y + z + w)))
+    c_obj, _ = convert_expressions(problem)
+
+    x_vals = np.array([1.0, 2.0])
+    y_vals = np.array([3.0, 4.0])
+    z_vals = np.array([5.0, 6.0])
+    w_vals = np.array([7.0, 8.0])
+    test_values = np.concatenate([x_vals, y_vals, z_vals, w_vals])
+
+    # Forward
+    result = diffengine.forward(c_obj, test_values)
+    expected = np.sum(x_vals + y_vals + z_vals + w_vals)
+    assert np.allclose(result, expected)
+
+    # TODO: Jacobian for sum(add(...)) patterns not yet supported
+
+
+def test_variable_used_multiple_times():
+    """Test sum(log(x) + exp(x) + x) - variable used 3+ times."""
+    x = cp.Variable(3)
+    problem = cp.Problem(cp.Minimize(cp.sum(cp.log(x) + cp.exp(x) + x)))
+    c_obj, _ = convert_expressions(problem)
+
+    test_values = np.array([1.0, 2.0, 3.0])
+
+    # Forward
+    result = diffengine.forward(c_obj, test_values)
+    expected = np.sum(np.log(test_values) + np.exp(test_values) + test_values)
+    assert np.allclose(result, expected)
+
+    # TODO: Jacobian for expressions with sum(variable) not yet supported
+
+
+def test_larger_variable():
+    """Test sum(log(x)) with larger variable (100 elements)."""
+    x = cp.Variable(100)
+    problem = cp.Problem(cp.Minimize(cp.sum(cp.log(x))))
+    c_obj, _ = convert_expressions(problem)
+
+    test_values = np.linspace(1.0, 10.0, 100)
+
+    # Forward
+    result = diffengine.forward(c_obj, test_values)
+    expected = np.sum(np.log(test_values))
+    assert np.allclose(result, expected)
+
+    # Jacobian
+    jac = get_jacobian(c_obj, test_values)
+    expected_jac = (1.0 / test_values).reshape(1, -1)
+    assert np.allclose(jac.toarray(), expected_jac)
+
+
+def test_matrix_variable():
+    """Test sum(log(X)) with 2D matrix variable (3x4)."""
+    X = cp.Variable((3, 4))
+    problem = cp.Problem(cp.Minimize(cp.sum(cp.log(X))))
+    c_obj, _ = convert_expressions(problem)
+
+    test_values = np.arange(1.0, 13.0)  # 12 elements
+
+    # Forward
+    result = diffengine.forward(c_obj, test_values)
+    expected = np.sum(np.log(test_values))
+    assert np.allclose(result, expected)
+
+    # Jacobian
+    jac = get_jacobian(c_obj, test_values)
+    expected_jac = (1.0 / test_values).reshape(1, -1)
+    assert np.allclose(jac.toarray(), expected_jac)
+
+
+def test_mixed_sizes():
+    """Test sum(log(a)) + sum(log(b)) + sum(log(c)) with different sized variables."""
+    a = cp.Variable(2)
+    b = cp.Variable(5)
+    c = cp.Variable(3)
+    problem = cp.Problem(cp.Minimize(cp.sum(cp.log(a)) + cp.sum(cp.log(b)) + cp.sum(cp.log(c))))
+    c_obj, _ = convert_expressions(problem)
+
+    a_vals = np.array([1.0, 2.0])
+    b_vals = np.array([1.0, 2.0, 3.0, 4.0, 5.0])
+    c_vals = np.array([1.0, 2.0, 3.0])
+    test_values = np.concatenate([a_vals, b_vals, c_vals])
+
+    # Forward
+    result = diffengine.forward(c_obj, test_values)
+    expected = np.sum(np.log(a_vals)) + np.sum(np.log(b_vals)) + np.sum(np.log(c_vals))
+    assert np.allclose(result, expected)
+
+    # Jacobian
+    jac = get_jacobian(c_obj, test_values)
+    expected_jac = (1.0 / test_values).reshape(1, -1)
+    assert np.allclose(jac.toarray(), expected_jac)
+
+
+def test_multiple_variables_log_exp():
+    """Test sum(log(a)) + sum(log(b)) + sum(exp(c)) + sum(exp(d))."""
+    a = cp.Variable(2)
+    b = cp.Variable(2)
+    c = cp.Variable(2)
+    d = cp.Variable(2)
+    obj = cp.sum(cp.log(a)) + cp.sum(cp.log(b)) + cp.sum(cp.exp(c)) + cp.sum(cp.exp(d))
+    problem = cp.Problem(cp.Minimize(obj))
+    c_obj, _ = convert_expressions(problem)
+
+    a_vals = np.array([1.0, 2.0])
+    b_vals = np.array([0.5, 1.0])
+    c_vals = np.array([0.1, 0.2])
+    d_vals = np.array([0.1, 0.1])
+    test_values = np.concatenate([a_vals, b_vals, c_vals, d_vals])
+
+    # Forward
+    result = diffengine.forward(c_obj, test_values)
+    expected = (np.sum(np.log(a_vals)) + np.sum(np.log(b_vals)) +
+                np.sum(np.exp(c_vals)) + np.sum(np.exp(d_vals)))
+    assert np.allclose(result, expected)
+
+    # Jacobian
+    jac = get_jacobian(c_obj, test_values)
+    df_da = 1.0 / a_vals
+    df_db = 1.0 / b_vals
+    df_dc = np.exp(c_vals)
+    df_dd = np.exp(d_vals)
+    expected_jac = np.concatenate([df_da, df_db, df_dc, df_dd]).reshape(1, -1)
+    assert np.allclose(jac.toarray(), expected_jac)
+
+
+def test_two_variables_separate_sums():
+    """Test sum(log(x) + log(y)) - two variables with separate elementwise ops."""
+    x = cp.Variable(2)
+    y = cp.Variable(2)
+    problem = cp.Problem(cp.Minimize(cp.sum(cp.log(x) + cp.log(y))))
+    c_obj, _ = convert_expressions(problem)
+
+    test_values = np.array([1.0, 2.0, 3.0, 4.0])
+
+    # Forward
+    result = diffengine.forward(c_obj, test_values)
+    expected = np.sum(np.log(test_values[:2]) + np.log(test_values[2:]))
+    assert np.allclose(result, expected)
+
+    # Jacobian
+    jac = get_jacobian(c_obj, test_values)
+    expected_jac = np.array([[1/1, 1/2, 1/3, 1/4]])
+    assert np.allclose(jac.toarray(), expected_jac)
+
+
+def test_complex_objective_elementwise_add():
+    """Test sum(log(x + y)) + sum(exp(y + z)) + sum(log(z + x)) - elementwise after add."""
+    x = cp.Variable(3)
+    y = cp.Variable(3)
+    z = cp.Variable(3)
+    obj = cp.sum(cp.log(x + y)) + cp.sum(cp.exp(y + z)) + cp.sum(cp.log(z + x))
+    problem = cp.Problem(cp.Minimize(obj))
+    c_obj, _ = convert_expressions(problem)
+
+    x_vals = np.array([1.0, 2.0, 3.0])
+    y_vals = np.array([0.5, 1.0, 1.5])
+    z_vals = np.array([0.2, 0.3, 0.4])
+    test_values = np.concatenate([x_vals, y_vals, z_vals])
+
+    # Forward
+    result = diffengine.forward(c_obj, test_values)
+    expected = (np.sum(np.log(x_vals + y_vals)) +
+                np.sum(np.exp(y_vals + z_vals)) +
+                np.sum(np.log(z_vals + x_vals)))
+    assert np.allclose(result, expected)
+
+    # TODO: Jacobian for elementwise(add(...)) patterns not yet supported
+
+
+def test_complex_objective_no_add():
+    """Test sum(log(x) + exp(y) + log(z)) - multiple elementwise ops without add composition."""
+    x = cp.Variable(2)
+    y = cp.Variable(2)
+    z = cp.Variable(2)
+    obj = cp.sum(cp.log(x) + cp.exp(y) + cp.log(z))
+    problem = cp.Problem(cp.Minimize(obj))
+    c_obj, _ = convert_expressions(problem)
+
+    x_vals = np.array([1.0, 2.0])
+    y_vals = np.array([0.5, 1.0])
+    z_vals = np.array([0.2, 0.3])
+    test_values = np.concatenate([x_vals, y_vals, z_vals])
+
+    # Forward
+    result = diffengine.forward(c_obj, test_values)
+    expected = np.sum(np.log(x_vals) + np.exp(y_vals) + np.log(z_vals))
+    assert np.allclose(result, expected)
+
+    # Jacobian
+    jac = get_jacobian(c_obj, test_values)
+    df_dx = 1.0 / x_vals
+    df_dy = np.exp(y_vals)
+    df_dz = 1.0 / z_vals
+    expected_jac = np.concatenate([df_dx, df_dy, df_dz]).reshape(1, -1)
+    assert np.allclose(jac.toarray(), expected_jac)
+
+
+def test_log_exp_identity():
+    """Test sum(log(exp(x))) = sum(x) identity - nested elementwise."""
+    x = cp.Variable(5)
+    problem = cp.Problem(cp.Minimize(cp.sum(cp.log(cp.exp(x)))))
+    c_obj, _ = convert_expressions(problem)
+
+    test_values = np.array([-1.0, 0.0, 1.0, 2.0, 3.0])
+
+    # Forward
+    result = diffengine.forward(c_obj, test_values)
+    expected = np.sum(test_values)  # log(exp(x)) = x
+    assert np.allclose(result, expected)
+
+    # TODO: Jacobian for nested elementwise compositions not yet supported
+
+
+if __name__ == "__main__":
+    test_sum_log()
+    test_sum_exp()
+    test_two_variables_elementwise_add()
+    test_variable_reuse()
+    test_four_variables_elementwise_add()
+    test_deep_nesting()
+    test_chained_additions()
+    test_variable_used_multiple_times()
+    test_larger_variable()
+    test_matrix_variable()
+    test_mixed_sizes()
+    test_multiple_variables_log_exp()
+    test_two_variables_separate_sums()
+    test_complex_objective_elementwise_add()
+    test_complex_objective_no_add()
+    test_log_exp_identity()
+    print("All tests passed!")
diff --git a/python/tests/test_problem_convert.py b/python/tests/test_problem_convert.py
new file mode 100644
index 0000000..1850fe2
--- /dev/null
+++ b/python/tests/test_problem_convert.py
@@ -0,0 +1,337 @@
+import cvxpy as cp
+import numpy as np
+from scipy import sparse
+import DNLP_diff_engine as diffengine
+from convert import Problem
+
+
+# ============ Low-level problem struct tests ============
+
+def test_problem_forward_lowlevel():
+    """Test problem_forward for objective and constraint values (low-level)."""
+    n_vars = 3
+    x_obj = diffengine.make_variable(n_vars, 1, 0, n_vars)
+    log_obj = diffengine.make_log(x_obj)
+    objective = diffengine.make_sum(log_obj, -1)
+
+    x_c = diffengine.make_variable(n_vars, 1, 0, n_vars)
+    log_c = diffengine.make_log(x_c)
+    constraints = [log_c]
+
+    prob = diffengine.make_problem(objective, constraints)
+    u = np.array([1.0, 2.0, 3.0])
+    diffengine.problem_allocate(prob, u)
+
+    obj_val, constraint_vals = diffengine.problem_forward(prob, u)
+
+    expected_obj = np.sum(np.log(u))
+    assert np.allclose(obj_val, expected_obj)
+    assert np.allclose(constraint_vals, np.log(u))
+
+
+def test_problem_gradient_lowlevel():
+    """Test problem_gradient for objective gradient (low-level)."""
+    n_vars = 3
+    x = diffengine.make_variable(n_vars, 1, 0, n_vars)
+    log_x = diffengine.make_log(x)
+    objective = diffengine.make_sum(log_x, -1)
+
+    prob = diffengine.make_problem(objective, [])
+    u = np.array([1.0, 2.0, 4.0])
+    diffengine.problem_allocate(prob, u)
+
+    grad = diffengine.problem_gradient(prob, u)
+    expected_grad = 1.0 / u
+    assert np.allclose(grad, expected_grad)
+
+
+def test_problem_jacobian_lowlevel():
+    """Test problem_jacobian for constraint jacobian (low-level)."""
+    n_vars = 2
+    x_obj = diffengine.make_variable(n_vars, 1, 0, n_vars)
+    log_obj = diffengine.make_log(x_obj)
+    objective = diffengine.make_sum(log_obj, -1)
+
+    x_c = diffengine.make_variable(n_vars, 1, 0, n_vars)
+    log_c = diffengine.make_log(x_c)
+    constraints = [log_c]
+
+    prob = diffengine.make_problem(objective, constraints)
+    u = np.array([2.0, 4.0])
+    diffengine.problem_allocate(prob, u)
+
+    data, indices, indptr, shape = diffengine.problem_jacobian(prob, u)
+    jac = sparse.csr_matrix((data, indices, indptr), shape=shape)
+
+    expected_jac = np.diag(1.0 / u)
+    assert np.allclose(jac.toarray(), expected_jac)
+
+
+def test_problem_no_constraints_lowlevel():
+    """Test Problem with no constraints (low-level)."""
+    n_vars = 3
+    x = diffengine.make_variable(n_vars, 1, 0, n_vars)
+    log_x = diffengine.make_log(x)
+    objective = diffengine.make_sum(log_x, -1)
+
+    prob = diffengine.make_problem(objective, [])
+    u = np.array([1.0, 2.0, 3.0])
+    diffengine.problem_allocate(prob, u)
+
+    obj_val, constraint_vals = diffengine.problem_forward(prob, u)
+    assert np.allclose(obj_val, np.sum(np.log(u)))
+    assert len(constraint_vals) == 0
+
+    grad = diffengine.problem_gradient(prob, u)
+    assert np.allclose(grad, 1.0 / u)
+
+    data, indices, indptr, shape = diffengine.problem_jacobian(prob, u)
+    jac = sparse.csr_matrix((data, indices, indptr), shape=shape)
+    assert jac.shape == (0, 3)
+
+
+# ============ Problem class tests using convert ============
+
+def test_problem_single_constraint():
+    """Test Problem class with single constraint using convert."""
+    x = cp.Variable(3)
+    obj = cp.sum(cp.log(x))
+    constraints = [cp.log(x)]  # log(x) as constraint expression
+
+    cvxpy_prob = cp.Problem(cp.Minimize(obj), constraints)
+    prob = Problem(cvxpy_prob)
+
+    u = np.array([1.0, 2.0, 3.0])
+    prob.allocate(u)
+
+    # Test forward
+    obj_val, constraint_vals = prob.forward(u)
+    assert np.allclose(obj_val, np.sum(np.log(u)))
+    assert np.allclose(constraint_vals, np.log(u))
+
+    # Test gradient
+    grad = prob.gradient(u)
+    assert np.allclose(grad, 1.0 / u)
+
+    # Test jacobian
+    jac = prob.jacobian(u)
+    expected_jac = np.diag(1.0 / u)
+    assert np.allclose(jac.toarray(), expected_jac)
+
+
+def test_problem_two_constraints():
+    """Test Problem class with two constraints."""
+    x = cp.Variable(2)
+    obj = cp.sum(cp.log(x))
+    constraints = [
+        cp.log(x),  # constraint 1: log(x), size 2
+        cp.exp(x),  # constraint 2: exp(x), size 2
+    ]
+
+    cvxpy_prob = cp.Problem(cp.Minimize(obj), constraints)
+    prob = Problem(cvxpy_prob)
+
+    u = np.array([1.0, 2.0])
+    prob.allocate(u)
+
+    # Test forward
+    obj_val, constraint_vals = prob.forward(u)
+    assert np.allclose(obj_val, np.sum(np.log(u)))
+    # Constraint values should be stacked: [log(1), log(2), exp(1), exp(2)]
+    expected_constraint_vals = np.concatenate([np.log(u), np.exp(u)])
+    assert np.allclose(constraint_vals, expected_constraint_vals)
+
+    # Test gradient
+    grad = prob.gradient(u)
+    assert np.allclose(grad, 1.0 / u)
+
+    # Test jacobian - stacked vertically
+    jac = prob.jacobian(u)
+    assert jac.shape == (4, 2)  # 2 constraints * 2 elements each
+    # First 2 rows: d(log(x))/dx = diag(1/x)
+    # Last 2 rows: d(exp(x))/dx = diag(exp(x))
+    expected_jac = np.vstack([np.diag(1.0 / u), np.diag(np.exp(u))])
+    assert np.allclose(jac.toarray(), expected_jac)
+
+
+def test_problem_three_constraints_different_sizes():
+    """Test Problem with three constraints of different types."""
+    x = cp.Variable(3)
+    obj = cp.sum(cp.exp(x))
+    constraints = [
+        cp.log(x),       # size 3
+        cp.exp(x),       # size 3
+        cp.sum(cp.log(x)),  # size 1 (scalar)
+    ]
+
+    cvxpy_prob = cp.Problem(cp.Minimize(obj), constraints)
+    prob = Problem(cvxpy_prob)
+
+    u = np.array([1.0, 2.0, 3.0])
+    prob.allocate(u)
+
+    # Test forward
+    obj_val, constraint_vals = prob.forward(u)
+    assert np.allclose(obj_val, np.sum(np.exp(u)))
+    # Constraint values: [log(x), exp(x), sum(log(x))]
+    expected_constraint_vals = np.concatenate([
+        np.log(u),
+        np.exp(u),
+        [np.sum(np.log(u))]
+    ])
+    assert np.allclose(constraint_vals, expected_constraint_vals)
+
+    # Test gradient of sum(exp(x))
+    grad = prob.gradient(u)
+    assert np.allclose(grad, np.exp(u))
+
+    # Test jacobian
+    jac = prob.jacobian(u)
+    assert jac.shape == (7, 3)  # 3 + 3 + 1 rows
+
+
+def test_problem_multiple_variables():
+    """Test Problem with multiple CVXPY variables."""
+    x = cp.Variable(2)
+    y = cp.Variable(2)
+    obj = cp.sum(cp.log(x)) + cp.sum(cp.exp(y))
+    constraints = [
+        cp.log(x),
+        cp.exp(y),
+    ]
+
+    cvxpy_prob = cp.Problem(cp.Minimize(obj), constraints)
+    prob = Problem(cvxpy_prob)
+
+    x_vals = np.array([1.0, 2.0])
+    y_vals = np.array([0.5, 1.0])
+    u = np.concatenate([x_vals, y_vals])
+    prob.allocate(u)
+
+    # Test forward
+    obj_val, constraint_vals = prob.forward(u)
+    expected_obj = np.sum(np.log(x_vals)) + np.sum(np.exp(y_vals))
+    assert np.allclose(obj_val, expected_obj)
+    expected_constraint_vals = np.concatenate([np.log(x_vals), np.exp(y_vals)])
+    assert np.allclose(constraint_vals, expected_constraint_vals)
+
+    # Test gradient
+    grad = prob.gradient(u)
+    expected_grad = np.concatenate([1.0 / x_vals, np.exp(y_vals)])
+    assert np.allclose(grad, expected_grad)
+
+    # Test jacobian
+    jac = prob.jacobian(u)
+    assert jac.shape == (4, 4)
+    # First constraint log(x) only depends on x (first 2 vars)
+    # Second constraint exp(y) only depends on y (last 2 vars)
+    expected_jac = np.zeros((4, 4))
+    expected_jac[0, 0] = 1.0 / x_vals[0]
+    expected_jac[1, 1] = 1.0 / x_vals[1]
+    expected_jac[2, 2] = np.exp(y_vals[0])
+    expected_jac[3, 3] = np.exp(y_vals[1])
+    assert np.allclose(jac.toarray(), expected_jac)
+
+
+def test_problem_no_constraints_convert():
+    """Test Problem class with no constraints using convert."""
+    x = cp.Variable(4)
+    obj = cp.sum(cp.log(x))
+
+    cvxpy_prob = cp.Problem(cp.Minimize(obj))
+    prob = Problem(cvxpy_prob)
+
+    u = np.array([1.0, 2.0, 3.0, 4.0])
+    prob.allocate(u)
+
+    obj_val, constraint_vals = prob.forward(u)
+    assert np.allclose(obj_val, np.sum(np.log(u)))
+    assert len(constraint_vals) == 0
+
+    grad = prob.gradient(u)
+    assert np.allclose(grad, 1.0 / u)
+
+    jac = prob.jacobian(u)
+    assert jac.shape == (0, 4)
+
+
+def test_problem_larger_scale():
+    """Test Problem with larger variables and multiple constraints."""
+    n = 50
+    x = cp.Variable(n)
+    obj = cp.sum(cp.log(x))
+    constraints = [
+        cp.log(x),
+        cp.exp(x),
+        cp.sum(cp.log(x)),
+        cp.sum(cp.exp(x)),
+    ]
+
+    cvxpy_prob = cp.Problem(cp.Minimize(obj), constraints)
+    prob = Problem(cvxpy_prob)
+
+    u = np.linspace(1.0, 5.0, n)
+    prob.allocate(u)
+
+    # Test forward
+    obj_val, constraint_vals = prob.forward(u)
+    assert np.allclose(obj_val, np.sum(np.log(u)))
+    expected_constraint_vals = np.concatenate([
+        np.log(u),
+        np.exp(u),
+        [np.sum(np.log(u))],
+        [np.sum(np.exp(u))],
+    ])
+    assert np.allclose(constraint_vals, expected_constraint_vals)
+
+    # Test gradient
+    grad = prob.gradient(u)
+    assert np.allclose(grad, 1.0 / u)
+
+    # Test jacobian shape
+    jac = prob.jacobian(u)
+    assert jac.shape == (n + n + 1 + 1, n)  # 102 x 50
+
+
+def test_problem_repeated_evaluations():
+    """Test Problem with repeated evaluations at different points."""
+    x = cp.Variable(3)
+    obj = cp.sum(cp.log(x))
+    constraints = [cp.exp(x)]
+
+    cvxpy_prob = cp.Problem(cp.Minimize(obj), constraints)
+    prob = Problem(cvxpy_prob)
+
+    u1 = np.array([1.0, 2.0, 3.0])
+    prob.allocate(u1)
+
+    # First evaluation
+    obj_val1, _ = prob.forward(u1)
+    grad1 = prob.gradient(u1)
+
+    # Second evaluation at different point
+    u2 = np.array([2.0, 3.0, 4.0])
+    obj_val2, _ = prob.forward(u2)
+    grad2 = prob.gradient(u2)
+
+    assert np.allclose(obj_val1, np.sum(np.log(u1)))
+    assert np.allclose(obj_val2, np.sum(np.log(u2)))
+    assert np.allclose(grad1, 1.0 / u1)
+    assert np.allclose(grad2, 1.0 / u2)
+
+
+if __name__ == "__main__":
+    # Low-level tests
+    test_problem_forward_lowlevel()
+    test_problem_gradient_lowlevel()
+    test_problem_jacobian_lowlevel()
+    test_problem_no_constraints_lowlevel()
+    # Problem class tests
+    test_problem_single_constraint()
+    test_problem_two_constraints()
+    test_problem_three_constraints_different_sizes()
+    test_problem_multiple_variables()
+    test_problem_no_constraints_convert()
+    test_problem_larger_scale()
+    test_problem_repeated_evaluations()
+    print("All problem tests passed!")
diff --git a/src/problem.c b/src/problem.c
new file mode 100644
index 0000000..37e06fc
--- /dev/null
+++ b/src/problem.c
@@ -0,0 +1,179 @@
+#include "problem.h"
+#include <stdlib.h>
+#include <string.h>
+
+problem *new_problem(expr *objective, expr **constraints, int n_constraints)
+{
+    problem *prob = (problem *) calloc(1, sizeof(problem));
+    if (!prob) return NULL;
+
+    prob->objective = objective;
+    expr_retain(objective);
+
+    /* Copy and retain constraints array */
+    prob->n_constraints = n_constraints;
+    if (n_constraints > 0)
+    {
+        prob->constraints = (expr **) malloc(n_constraints * sizeof(expr *));
+        for (int i = 0; i < n_constraints; i++)
+        {
+            prob->constraints[i] = constraints[i];
+            expr_retain(constraints[i]);
+        }
+    }
+    else
+    {
+        prob->constraints = NULL;
+    }
+
+    /* Compute total constraint size */
+    prob->total_constraint_size = 0;
+    for (int i = 0; i < n_constraints; i++)
+    {
+        prob->total_constraint_size += constraints[i]->size;
+    }
+
+    prob->n_vars = objective->n_vars;
+
+    /* Initialize allocated pointers to NULL */
+    prob->constraint_values = NULL;
+    prob->gradient_values = NULL;
+    prob->stacked_jac = NULL;
+
+    return prob;
+}
+
+void problem_allocate(problem *prob, const double *u)
+{
+    /* 1. Allocate constraint values array */
+    if (prob->total_constraint_size > 0)
+    {
+        prob->constraint_values = (double *) calloc(prob->total_constraint_size, sizeof(double));
+    }
+
+    /* 2. Allocate gradient values array */
+    prob->gradient_values = (double *) calloc(prob->n_vars, sizeof(double));
+
+    /* 3. Initialize objective jacobian */
+    prob->objective->forward(prob->objective, u);
+    prob->objective->jacobian_init(prob->objective);
+
+    /* 4. Initialize constraint jacobians and count total nnz */
+    int total_nnz = 0;
+    for (int i = 0; i < prob->n_constraints; i++)
+    {
+        expr *c = prob->constraints[i];
+        c->forward(c, u);
+        c->jacobian_init(c);
+        total_nnz += c->jacobian->nnz;
+    }
+
+    /* 5. Allocate stacked jacobian */
+    if (prob->total_constraint_size > 0)
+    {
+        prob->stacked_jac = new_csr_matrix(prob->total_constraint_size, prob->n_vars, total_nnz);
+    }
+}
+
+void free_problem(problem *prob)
+{
+    if (prob == NULL) return;
+
+    /* Free allocated arrays */
+    free(prob->constraint_values);
+    free(prob->gradient_values);
+    free_csr_matrix(prob->stacked_jac);
+
+    /* Free expressions (decrements refcount) */
+    free_expr(prob->objective);
+    for (int i = 0; i < prob->n_constraints; i++)
+    {
+        free_expr(prob->constraints[i]);
+    }
+    free(prob->constraints);
+
+    /* Free problem struct */
+    free(prob);
+}
+
+double problem_forward(problem *prob, const double *u)
+{
+    /* Evaluate objective */
+    prob->objective->forward(prob->objective, u);
+    double obj_val = prob->objective->value[0];
+
+    /* Evaluate constraints and copy values */
+    int offset = 0;
+    for (int i = 0; i < prob->n_constraints; i++)
+    {
+        expr *c = prob->constraints[i];
+        c->forward(c, u);
+        memcpy(prob->constraint_values + offset, c->value, c->size * sizeof(double));
+        offset += c->size;
+    }
+
+    return obj_val;
+}
+
+double *problem_gradient(problem *prob, const double *u)
+{
+    /* Forward and jacobian on objective */
+    prob->objective->forward(prob->objective, u);
+    prob->objective->eval_jacobian(prob->objective);
+
+    /* Zero gradient array */
+    memset(prob->gradient_values, 0, prob->n_vars * sizeof(double));
+
+    /* Copy sparse jacobian row to dense gradient
+     * Objective jacobian is 1 x n_vars */
+    CSR_Matrix *jac = prob->objective->jacobian;
+    for (int k = jac->p[0]; k < jac->p[1]; k++)
+    {
+        int col = jac->i[k];
+        prob->gradient_values[col] = jac->x[k];
+    }
+
+    return prob->gradient_values;
+}
+
+CSR_Matrix *problem_jacobian(problem *prob, const double *u)
+{
+    CSR_Matrix *stacked = prob->stacked_jac;
+
+    /* Initialize row pointers */
+    stacked->p[0] = 0;
+
+    int row_offset = 0;
+    int nnz_offset = 0;
+
+    for (int i = 0; i < prob->n_constraints; i++)
+    {
+        expr *c = prob->constraints[i];
+
+        /* Forward and eval jacobian */
+        c->forward(c, u);
+        c->eval_jacobian(c);
+
+        CSR_Matrix *cjac = c->jacobian;
+
+        /* Copy row pointers with offset */
+        for (int r = 0; r < cjac->m; r++)
+        {
+            int row_nnz = cjac->p[r + 1] - cjac->p[r];
+            stacked->p[row_offset + r + 1] = stacked->p[row_offset + r] + row_nnz;
+        }
+
+        /* Copy column indices and values */
+        int constraint_nnz = cjac->p[cjac->m];
+        memcpy(stacked->i + nnz_offset, cjac->i, constraint_nnz * sizeof(int));
+        memcpy(stacked->x + nnz_offset, cjac->x, constraint_nnz * sizeof(double));
+
+        row_offset += cjac->m;
+        nnz_offset += constraint_nnz;
+    }
+
+    /* Update actual nnz (may be less than allocated) */
+    stacked->nnz = nnz_offset;
+
+    return stacked;
+}
diff --git a/tests/problem/test_problem.h b/tests/problem/test_problem.h
new file mode 100644
index 0000000..47d3e97
--- /dev/null
+++ b/tests/problem/test_problem.h
@@ -0,0 +1,162 @@
+#ifndef TEST_PROBLEM_H
+#define TEST_PROBLEM_H
+
+#include <math.h>
+#include <stdio.h>
+
+#include "affine.h"
+#include "elementwise_univariate.h"
+#include "expr.h"
+#include "minunit.h"
+#include "problem.h"
+#include "test_helpers.h"
+
+/*
+ * Test problem: minimize sum(log(x))
+ *               subject to x >= 1 (as x - 1 >= 0)
+ *
+ * With x of size 3, n_vars = 3
+ */
+const char *test_problem_new_free(void)
+{
+    /* Create expressions */
+    expr *x = new_variable(3, 1, 0, 3);
+    expr *log_x = new_log(x);
+    expr *objective = new_sum(log_x, -1);
+
+    /* Create constraint: x - 1 (represented as just x for simplicity) */
+    expr *x_constraint = new_variable(3, 1, 0, 3);
+
+    expr *constraints[1] = {x_constraint};
+
+    /* Create problem */
+    problem *prob = new_problem(objective, constraints, 1);
+
+    mu_assert("new_problem failed", prob != NULL);
+    mu_assert("n_vars wrong", prob->n_vars == 3);
+    mu_assert("n_constraints wrong", prob->n_constraints == 1);
+    mu_assert("total_constraint_size wrong", prob->total_constraint_size == 3);
+
+    /* Free problem (also frees expressions via refcount) */
+    free_problem(prob);
+
+    /* Free original expression refs */
+    free_expr(objective);
+    free_expr(x_constraint);
+
+    return 0;
+}
+
+/*
+ * Test problem_forward: minimize sum(log(x))
+ *                       subject to x (as constraint)
+ */
+const char *test_problem_forward(void)
+{
+    expr *x = new_variable(3, 1, 0, 3);
+    expr *log_x = new_log(x);
+    expr *objective = new_sum(log_x, -1);
+
+    expr *x_constraint = new_variable(3, 1, 0, 3);
+    expr *constraints[1] = {x_constraint};
+
+    problem *prob = new_problem(objective, constraints, 1);
+
+    double u[3] = {1.0, 2.0, 3.0};
+    problem_allocate(prob, u);
+
+    double obj_val = problem_forward(prob, u);
+
+    /* Expected: sum(log([1, 2, 3])) = 0 + log(2) + log(3) */
+    double expected_obj = log(1.0) + log(2.0) + log(3.0);
+    mu_assert("objective value wrong", fabs(obj_val - expected_obj) < 1e-10);
+
+    /* Constraint values should be [1, 2, 3] */
+    mu_assert("constraint[0] wrong", fabs(prob->constraint_values[0] - 1.0) < 1e-10);
+    mu_assert("constraint[1] wrong", fabs(prob->constraint_values[1] - 2.0) < 1e-10);
+    mu_assert("constraint[2] wrong", fabs(prob->constraint_values[2] - 3.0) < 1e-10);
+
+    free_problem(prob);
+    free_expr(objective);
+    free_expr(x_constraint);
+
+    return 0;
+}
+
+/*
+ * Test problem_gradient: gradient of sum(log(x)) = [1/x1, 1/x2, 1/x3]
+ */
+const char *test_problem_gradient(void)
+{
+    expr *x = new_variable(3, 1, 0, 3);
+    expr *log_x = new_log(x);
+    expr *objective = new_sum(log_x, -1);
+
+    problem *prob = new_problem(objective, NULL, 0);
+
+    double u[3] = {1.0, 2.0, 4.0};
+    problem_allocate(prob, u);
+
+    double *grad = problem_gradient(prob, u);
+
+    /* Expected gradient: [1/1, 1/2, 1/4] = [1.0, 0.5, 0.25] */
+    mu_assert("grad[0] wrong", fabs(grad[0] - 1.0) < 1e-10);
+    mu_assert("grad[1] wrong", fabs(grad[1] - 0.5) < 1e-10);
+    mu_assert("grad[2] wrong", fabs(grad[2] - 0.25) < 1e-10);
+
+    free_problem(prob);
+    free_expr(objective);
+
+    return 0;
+}
+
+/*
+ * Test problem_jacobian: one constraint log(x)
+ * Jacobian of log(x): diag([1/x1, 1/x2])
+ */
+const char *test_problem_jacobian(void)
+{
+    int n_vars = 2;
+
+    /* Create separate expression trees */
+    expr *x_obj = new_variable(2, 1, 0, n_vars);
+    expr *log_obj = new_log(x_obj);
+    expr *objective = new_sum(log_obj, -1);
+
+    expr *x_c1 = new_variable(2, 1, 0, n_vars);
+    expr *log_c1 = new_log(x_c1);
+
+    expr *constraints[1] = {log_c1};
+
+    problem *prob = new_problem(objective, constraints, 1);
+
+    double u[2] = {2.0, 4.0};
+    problem_allocate(prob, u);
+
+    CSR_Matrix *jac = problem_jacobian(prob, u);
+
+    /* Check dimensions */
+    mu_assert("jac rows wrong", jac->m == 2);
+    mu_assert("jac cols wrong", jac->n == 2);
+
+    /* Check row pointers: each row has 1 element */
+    mu_assert("jac->p[0] wrong", jac->p[0] == 0);
+    mu_assert("jac->p[1] wrong", jac->p[1] == 1);
+    mu_assert("jac->p[2] wrong", jac->p[2] == 2);
+
+    /* Check column indices */
+    mu_assert("jac->i[0] wrong", jac->i[0] == 0);
+    mu_assert("jac->i[1] wrong", jac->i[1] == 1);
+
+    /* Check values: [1/2, 1/4] */
+    mu_assert("jac->x[0] wrong", fabs(jac->x[0] - 0.5) < 1e-10);
+    mu_assert("jac->x[1] wrong", fabs(jac->x[1] - 0.25) < 1e-10);
+
+    free_problem(prob);
+    free_expr(objective);
+    free_expr(log_c1);
+
+    return 0;
+}
+
+#endif /* TEST_PROBLEM_H */

From d1ae7d551cc16cae7a9b6b73440ac099553070aa Mon Sep 17 00:00:00 2001
From: William Zijie Zhang <william@gridmatic.com>
Date: Mon, 5 Jan 2026 00:28:02 -0500
Subject: [PATCH 03/27] fix freeing of expressions in problem

---
 include/problem.h            |  1 +
 python/convert.py            |  6 ------
 src/problem.c                |  5 ++---
 tests/problem/test_problem.h | 11 +----------
 4 files changed, 4 insertions(+), 19 deletions(-)

diff --git a/include/problem.h b/include/problem.h
index 97b0e9f..956f731 100644
--- a/include/problem.h
+++ b/include/problem.h
@@ -18,6 +18,7 @@ typedef struct problem
     CSR_Matrix *stacked_jac;
 } problem;
 
+/* Takes ownership of objective and constraints - caller should not free them */
 problem *new_problem(expr *objective, expr **constraints, int n_constraints);
 void problem_allocate(problem *prob, const double *u);
 void free_problem(problem *prob);
diff --git a/python/convert.py b/python/convert.py
index bf8ccd4..faeb0ba 100644
--- a/python/convert.py
+++ b/python/convert.py
@@ -5,12 +5,6 @@
 import DNLP_diff_engine as diffengine
 
 
-def get_jacobian(c_expr, values):
-    """Compute jacobian and return as scipy sparse CSR matrix."""
-    data, indices, indptr, shape = diffengine.jacobian(c_expr, values)
-    return sparse.csr_matrix((data, indices, indptr), shape=shape)
-
-
 def _chain_add(children):
     """Chain multiple children with binary adds: a + b + c -> add(add(a, b), c)."""
     result = children[0]
diff --git a/src/problem.c b/src/problem.c
index 37e06fc..249569a 100644
--- a/src/problem.c
+++ b/src/problem.c
@@ -7,10 +7,10 @@ problem *new_problem(expr *objective, expr **constraints, int n_constraints)
     problem *prob = (problem *) calloc(1, sizeof(problem));
     if (!prob) return NULL;
 
+    /* Take ownership of objective (no retain - caller transfers ownership) */
     prob->objective = objective;
-    expr_retain(objective);
 
-    /* Copy and retain constraints array */
+    /* Copy constraints array (take ownership, no retain) */
     prob->n_constraints = n_constraints;
     if (n_constraints > 0)
     {
@@ -18,7 +18,6 @@ problem *new_problem(expr *objective, expr **constraints, int n_constraints)
         for (int i = 0; i < n_constraints; i++)
         {
             prob->constraints[i] = constraints[i];
-            expr_retain(constraints[i]);
         }
     }
     else
diff --git a/tests/problem/test_problem.h b/tests/problem/test_problem.h
index 47d3e97..8349144 100644
--- a/tests/problem/test_problem.h
+++ b/tests/problem/test_problem.h
@@ -37,13 +37,9 @@ const char *test_problem_new_free(void)
     mu_assert("n_constraints wrong", prob->n_constraints == 1);
     mu_assert("total_constraint_size wrong", prob->total_constraint_size == 3);
 
-    /* Free problem (also frees expressions via refcount) */
+    /* Free problem (owns and frees all expressions) */
     free_problem(prob);
 
-    /* Free original expression refs */
-    free_expr(objective);
-    free_expr(x_constraint);
-
     return 0;
 }
 
@@ -77,8 +73,6 @@ const char *test_problem_forward(void)
     mu_assert("constraint[2] wrong", fabs(prob->constraint_values[2] - 3.0) < 1e-10);
 
     free_problem(prob);
-    free_expr(objective);
-    free_expr(x_constraint);
 
     return 0;
 }
@@ -105,7 +99,6 @@ const char *test_problem_gradient(void)
     mu_assert("grad[2] wrong", fabs(grad[2] - 0.25) < 1e-10);
 
     free_problem(prob);
-    free_expr(objective);
 
     return 0;
 }
@@ -153,8 +146,6 @@ const char *test_problem_jacobian(void)
     mu_assert("jac->x[1] wrong", fabs(jac->x[1] - 0.25) < 1e-10);
 
     free_problem(prob);
-    free_expr(objective);
-    free_expr(log_c1);
 
     return 0;
 }

From e12c7522a3873f040efef49237425a95a9ac4762 Mon Sep 17 00:00:00 2001
From: William Zijie Zhang <william@gridmatic.com>
Date: Mon, 5 Jan 2026 00:40:49 -0500
Subject: [PATCH 04/27] free expressions recursively in problem

---
 src/problem.c | 14 +++++++++++++-
 1 file changed, 13 insertions(+), 1 deletion(-)

diff --git a/src/problem.c b/src/problem.c
index 249569a..3c49cd6 100644
--- a/src/problem.c
+++ b/src/problem.c
@@ -2,6 +2,16 @@
 #include <stdlib.h>
 #include <string.h>
 
+/* Release intermediate refs in expression tree (call free_expr on all children) */
+static void release_child_refs(expr *node)
+{
+    if (node == NULL) return;
+    release_child_refs(node->left);
+    release_child_refs(node->right);
+    free_expr(node->left);
+    free_expr(node->right);
+}
+
 problem *new_problem(expr *objective, expr **constraints, int n_constraints)
 {
     problem *prob = (problem *) calloc(1, sizeof(problem));
@@ -83,10 +93,12 @@ void free_problem(problem *prob)
     free(prob->gradient_values);
     free_csr_matrix(prob->stacked_jac);
 
-    /* Free expressions (decrements refcount) */
+    /* Free expression trees: release intermediate refs first, then free root */
+    release_child_refs(prob->objective);
     free_expr(prob->objective);
     for (int i = 0; i < prob->n_constraints; i++)
     {
+        release_child_refs(prob->constraints[i]);
         free_expr(prob->constraints[i]);
     }
     free(prob->constraints);

From cca4fbc363eee7b6e67a066d9b2213d230a4e1e0 Mon Sep 17 00:00:00 2001
From: William Zijie Zhang <william@gridmatic.com>
Date: Mon, 5 Jan 2026 13:58:53 -0500
Subject: [PATCH 05/27] adds shared variables and freeing them using a visited
 node tracker

---
 src/problem.c                | 77 ++++++++++++++++++++++++++++--------
 tests/all_tests.c            |  1 +
 tests/problem/test_problem.h | 75 +++++++++++++++++++++++++++++++++++
 3 files changed, 137 insertions(+), 16 deletions(-)

diff --git a/src/problem.c b/src/problem.c
index 3c49cd6..b3ee464 100644
--- a/src/problem.c
+++ b/src/problem.c
@@ -2,14 +2,60 @@
 #include <stdlib.h>
 #include <string.h>
 
-/* Release intermediate refs in expression tree (call free_expr on all children) */
-static void release_child_refs(expr *node)
+/* Simple visited set for tracking freed nodes (handles up to 1024 unique nodes) */
+#define MAX_VISITED 1024
+
+typedef struct
+{
+    expr *nodes[MAX_VISITED];
+    int count;
+} VisitedSet;
+
+static void visited_init(VisitedSet *v)
 {
-    if (node == NULL) return;
-    release_child_refs(node->left);
-    release_child_refs(node->right);
-    free_expr(node->left);
-    free_expr(node->right);
+    v->count = 0;
+}
+
+static int visited_contains(VisitedSet *v, expr *node)
+{
+    for (int i = 0; i < v->count; i++)
+    {
+        if (v->nodes[i] == node) return 1;
+    }
+    return 0;
+}
+
+static void visited_add(VisitedSet *v, expr *node)
+{
+    if (v->count < MAX_VISITED)
+    {
+        v->nodes[v->count++] = node;
+    }
+}
+
+/* Release refs and free nodes, tracking visited to handle sharing */
+static void free_expr_tree_visited(expr *node, VisitedSet *visited)
+{
+    if (node == NULL || visited_contains(visited, node)) return;
+    visited_add(visited, node);
+
+    /* Recursively process children first */
+    free_expr_tree_visited(node->left, visited);
+    free_expr_tree_visited(node->right, visited);
+
+    /* Free this node's resources */
+    free(node->value);
+    free_csr_matrix(node->jacobian);
+    free_csr_matrix(node->wsum_hess);
+    free(node->dwork);
+    free(node->iwork);
+
+    if (node->free_type_data)
+    {
+        node->free_type_data(node);
+    }
+
+    free(node);
 }
 
 problem *new_problem(expr *objective, expr **constraints, int n_constraints)
@@ -93,13 +139,14 @@ void free_problem(problem *prob)
     free(prob->gradient_values);
     free_csr_matrix(prob->stacked_jac);
 
-    /* Free expression trees: release intermediate refs first, then free root */
-    release_child_refs(prob->objective);
-    free_expr(prob->objective);
+    /* Free expression trees with shared visited set to handle node sharing */
+    VisitedSet visited;
+    visited_init(&visited);
+
+    free_expr_tree_visited(prob->objective, &visited);
     for (int i = 0; i < prob->n_constraints; i++)
     {
-        release_child_refs(prob->constraints[i]);
-        free_expr(prob->constraints[i]);
+        free_expr_tree_visited(prob->constraints[i], &visited);
     }
     free(prob->constraints);
 
@@ -128,8 +175,7 @@ double problem_forward(problem *prob, const double *u)
 
 double *problem_gradient(problem *prob, const double *u)
 {
-    /* Forward and jacobian on objective */
-    prob->objective->forward(prob->objective, u);
+    /* Jacobian on objective */
     prob->objective->eval_jacobian(prob->objective);
 
     /* Zero gradient array */
@@ -161,8 +207,7 @@ CSR_Matrix *problem_jacobian(problem *prob, const double *u)
     {
         expr *c = prob->constraints[i];
 
-        /* Forward and eval jacobian */
-        c->forward(c, u);
+        /* Evaluate jacobian */
         c->eval_jacobian(c);
 
         CSR_Matrix *cjac = c->jacobian;
diff --git a/tests/all_tests.c b/tests/all_tests.c
index a7cc644..7bc8aad 100644
--- a/tests/all_tests.c
+++ b/tests/all_tests.c
@@ -124,6 +124,7 @@ int main(void)
     mu_run_test(test_problem_forward, tests_run);
     mu_run_test(test_problem_gradient, tests_run);
     mu_run_test(test_problem_jacobian, tests_run);
+    mu_run_test(test_problem_jacobian_multi, tests_run);
 
     printf("\n=== All %d tests passed ===\n", tests_run);
 
diff --git a/tests/problem/test_problem.h b/tests/problem/test_problem.h
index 8349144..e3115d1 100644
--- a/tests/problem/test_problem.h
+++ b/tests/problem/test_problem.h
@@ -150,4 +150,79 @@ const char *test_problem_jacobian(void)
     return 0;
 }
 
+/*
+ * Test problem_jacobian with multiple constraints and SHARED variable:
+ *   Constraint 1: log(x) -> Jacobian diag([1/x1, 1/x2])
+ *   Constraint 2: exp(x) -> Jacobian diag([exp(x1), exp(x2)])
+ *
+ * Stacked jacobian (4x2):
+ *   [ 1/x1    0    ]
+ *   [  0    1/x2   ]
+ *   [exp(x1)  0    ]
+ *   [  0    exp(x2)]
+ *
+ * Note: All expressions share the same variable node x, testing that
+ * free_problem correctly handles shared nodes without double-free.
+ */
+const char *test_problem_jacobian_multi(void)
+{
+    int n_vars = 2;
+
+    /* Single shared variable used across all expressions */
+    expr *x = new_variable(2, 1, 0, n_vars);
+
+    /* Objective: sum(log(x)) */
+    expr *log_obj = new_log(x);
+    expr *objective = new_sum(log_obj, -1);
+
+    /* Constraint 1: log(x) - shares x */
+    expr *log_c1 = new_log(x);
+
+    /* Constraint 2: exp(x) - shares x */
+    expr *exp_c2 = new_exp(x);
+
+    expr *constraints[2] = {log_c1, exp_c2};
+
+    problem *prob = new_problem(objective, constraints, 2);
+
+    double u[2] = {2.0, 4.0};
+    problem_allocate(prob, u);
+    problem_forward(prob, u);
+
+    CSR_Matrix *jac = problem_jacobian(prob, u);
+
+    /* Check dimensions: 4 rows (2 + 2), 2 cols */
+    mu_assert("jac rows wrong", jac->m == 4);
+    mu_assert("jac cols wrong", jac->n == 2);
+    mu_assert("jac nnz wrong", jac->nnz == 4);
+
+    /* Check row pointers: each row has 1 element */
+    mu_assert("jac->p[0] wrong", jac->p[0] == 0);
+    mu_assert("jac->p[1] wrong", jac->p[1] == 1);
+    mu_assert("jac->p[2] wrong", jac->p[2] == 2);
+    mu_assert("jac->p[3] wrong", jac->p[3] == 3);
+    mu_assert("jac->p[4] wrong", jac->p[4] == 4);
+
+    /* Check column indices: diagonal pattern */
+    mu_assert("jac->i[0] wrong", jac->i[0] == 0);
+    mu_assert("jac->i[1] wrong", jac->i[1] == 1);
+    mu_assert("jac->i[2] wrong", jac->i[2] == 0);
+    mu_assert("jac->i[3] wrong", jac->i[3] == 1);
+
+    /* Check values:
+     * Row 0: 1/2 = 0.5
+     * Row 1: 1/4 = 0.25
+     * Row 2: exp(2) ≈ 7.389
+     * Row 3: exp(4) ≈ 54.598
+     */
+    mu_assert("jac->x[0] wrong", fabs(jac->x[0] - 0.5) < 1e-10);
+    mu_assert("jac->x[1] wrong", fabs(jac->x[1] - 0.25) < 1e-10);
+    mu_assert("jac->x[2] wrong", fabs(jac->x[2] - exp(2.0)) < 1e-10);
+    mu_assert("jac->x[3] wrong", fabs(jac->x[3] - exp(4.0)) < 1e-10);
+
+    free_problem(prob);
+
+    return 0;
+}
+
 #endif /* TEST_PROBLEM_H */

From 8ea75ab8be5113db0cfdb1105ae43de767a2fe1c Mon Sep 17 00:00:00 2001
From: William Zijie Zhang <william@gridmatic.com>
Date: Mon, 5 Jan 2026 15:26:48 -0500
Subject: [PATCH 06/27] adds constraint forward as well

---
 include/problem.h                    |  3 +-
 python/bindings.c                    | 51 ++++++++++++++++++++
 python/convert.py                    |  4 ++
 python/tests/test_problem_convert.py | 50 +++++++++++++++++++
 src/problem.c                        | 30 ++++++++----
 tests/all_tests.c                    |  1 +
 tests/problem/test_problem.h         | 72 +++++++++++++++++++++++++++-
 7 files changed, 201 insertions(+), 10 deletions(-)

diff --git a/include/problem.h b/include/problem.h
index 956f731..9c29399 100644
--- a/include/problem.h
+++ b/include/problem.h
@@ -18,12 +18,13 @@ typedef struct problem
     CSR_Matrix *stacked_jac;
 } problem;
 
-/* Takes ownership of objective and constraints - caller should not free them */
+/* Retains objective and constraints (shared ownership with caller) */
 problem *new_problem(expr *objective, expr **constraints, int n_constraints);
 void problem_allocate(problem *prob, const double *u);
 void free_problem(problem *prob);
 
 double problem_forward(problem *prob, const double *u);
+double *problem_constraint_forward(problem *prob, const double *u);
 double *problem_gradient(problem *prob, const double *u);
 CSR_Matrix *problem_jacobian(problem *prob, const double *u);
 
diff --git a/python/bindings.c b/python/bindings.c
index dce1477..06354a3 100644
--- a/python/bindings.c
+++ b/python/bindings.c
@@ -388,6 +388,56 @@ static PyObject *py_problem_forward(PyObject *self, PyObject *args)
     return Py_BuildValue("(dO)", obj_val, constraint_vals);
 }
 
+static PyObject *py_problem_constraint_forward(PyObject *self, PyObject *args)
+{
+    PyObject *prob_capsule;
+    PyObject *u_obj;
+    if (!PyArg_ParseTuple(args, "OO", &prob_capsule, &u_obj))
+    {
+        return NULL;
+    }
+
+    problem *prob =
+        (problem *) PyCapsule_GetPointer(prob_capsule, PROBLEM_CAPSULE_NAME);
+    if (!prob)
+    {
+        PyErr_SetString(PyExc_ValueError, "invalid problem capsule");
+        return NULL;
+    }
+
+    PyArrayObject *u_array =
+        (PyArrayObject *) PyArray_FROM_OTF(u_obj, NPY_DOUBLE, NPY_ARRAY_IN_ARRAY);
+    if (!u_array)
+    {
+        return NULL;
+    }
+
+    double *constraint_vals =
+        problem_constraint_forward(prob, (const double *) PyArray_DATA(u_array));
+
+    PyObject *out = NULL;
+    if (prob->total_constraint_size > 0)
+    {
+        npy_intp size = prob->total_constraint_size;
+        out = PyArray_SimpleNew(1, &size, NPY_DOUBLE);
+        if (!out)
+        {
+            Py_DECREF(u_array);
+            return NULL;
+        }
+        memcpy(PyArray_DATA((PyArrayObject *) out), constraint_vals,
+               size * sizeof(double));
+    }
+    else
+    {
+        npy_intp size = 0;
+        out = PyArray_SimpleNew(1, &size, NPY_DOUBLE);
+    }
+
+    Py_DECREF(u_array);
+    return out;
+}
+
 static PyObject *py_problem_gradient(PyObject *self, PyObject *args)
 {
     PyObject *prob_capsule;
@@ -500,6 +550,7 @@ static PyMethodDef DNLPMethods[] = {
     {"make_problem", py_make_problem, METH_VARARGS, "Create problem from objective and constraints"},
     {"problem_allocate", py_problem_allocate, METH_VARARGS, "Allocate problem resources"},
     {"problem_forward", py_problem_forward, METH_VARARGS, "Evaluate objective and constraints"},
+    {"problem_constraint_forward", py_problem_constraint_forward, METH_VARARGS, "Evaluate constraints only"},
     {"problem_gradient", py_problem_gradient, METH_VARARGS, "Compute objective gradient"},
     {"problem_jacobian", py_problem_jacobian, METH_VARARGS, "Compute constraint jacobian"},
     {NULL, NULL, 0, NULL}};
diff --git a/python/convert.py b/python/convert.py
index faeb0ba..900a32d 100644
--- a/python/convert.py
+++ b/python/convert.py
@@ -132,6 +132,10 @@ def forward(self, u: np.ndarray) -> tuple[float, np.ndarray]:
         """Evaluate objective and constraints. Returns (obj_value, constraint_values)."""
         return diffengine.problem_forward(self._capsule, u)
 
+    def constraint_forward(self, u: np.ndarray) -> np.ndarray:
+        """Evaluate constraints only. Returns constraint_values array."""
+        return diffengine.problem_constraint_forward(self._capsule, u)
+
     def gradient(self, u: np.ndarray) -> np.ndarray:
         """Compute gradient of objective. Returns gradient array."""
         return diffengine.problem_gradient(self._capsule, u)
diff --git a/python/tests/test_problem_convert.py b/python/tests/test_problem_convert.py
index 1850fe2..de605cd 100644
--- a/python/tests/test_problem_convert.py
+++ b/python/tests/test_problem_convert.py
@@ -67,6 +67,29 @@ def test_problem_jacobian_lowlevel():
     assert np.allclose(jac.toarray(), expected_jac)
 
 
+def test_problem_constraint_forward_lowlevel():
+    """Test problem_constraint_forward for constraint values only (low-level)."""
+    n_vars = 2
+    x = diffengine.make_variable(n_vars, 1, 0, n_vars)
+
+    log_obj = diffengine.make_log(x)
+    objective = diffengine.make_sum(log_obj, -1)
+
+    log_c = diffengine.make_log(x)
+    exp_c = diffengine.make_exp(x)
+    constraints = [log_c, exp_c]
+
+    prob = diffengine.make_problem(objective, constraints)
+    u = np.array([2.0, 4.0])
+    diffengine.problem_allocate(prob, u)
+
+    constraint_vals = diffengine.problem_constraint_forward(prob, u)
+
+    # Expected: [log(2), log(4), exp(2), exp(4)]
+    expected = np.concatenate([np.log(u), np.exp(u)])
+    assert np.allclose(constraint_vals, expected)
+
+
 def test_problem_no_constraints_lowlevel():
     """Test Problem with no constraints (low-level)."""
     n_vars = 3
@@ -320,11 +343,37 @@ def test_problem_repeated_evaluations():
     assert np.allclose(grad2, 1.0 / u2)
 
 
+def test_problem_constraint_forward():
+    """Test Problem.constraint_forward for constraint values only."""
+    x = cp.Variable(2)
+    obj = cp.sum(cp.log(x))
+    constraints = [
+        cp.log(x),
+        cp.exp(x),
+    ]
+
+    cvxpy_prob = cp.Problem(cp.Minimize(obj), constraints)
+    prob = Problem(cvxpy_prob)
+
+    u = np.array([2.0, 4.0])
+    prob.allocate(u)
+
+    # Test constraint_forward
+    constraint_vals = prob.constraint_forward(u)
+    expected = np.concatenate([np.log(u), np.exp(u)])
+    assert np.allclose(constraint_vals, expected)
+
+    # Verify it gives same result as forward's constraint values
+    _, forward_constraint_vals = prob.forward(u)
+    assert np.allclose(constraint_vals, forward_constraint_vals)
+
+
 if __name__ == "__main__":
     # Low-level tests
     test_problem_forward_lowlevel()
     test_problem_gradient_lowlevel()
     test_problem_jacobian_lowlevel()
+    test_problem_constraint_forward_lowlevel()
     test_problem_no_constraints_lowlevel()
     # Problem class tests
     test_problem_single_constraint()
@@ -334,4 +383,5 @@ def test_problem_repeated_evaluations():
     test_problem_no_constraints_convert()
     test_problem_larger_scale()
     test_problem_repeated_evaluations()
+    test_problem_constraint_forward()
     print("All problem tests passed!")
diff --git a/src/problem.c b/src/problem.c
index b3ee464..45f0707 100644
--- a/src/problem.c
+++ b/src/problem.c
@@ -63,10 +63,11 @@ problem *new_problem(expr *objective, expr **constraints, int n_constraints)
     problem *prob = (problem *) calloc(1, sizeof(problem));
     if (!prob) return NULL;
 
-    /* Take ownership of objective (no retain - caller transfers ownership) */
+    /* Retain objective (shared ownership with caller) */
     prob->objective = objective;
+    expr_retain(objective);
 
-    /* Copy constraints array (take ownership, no retain) */
+    /* Copy and retain constraints array */
     prob->n_constraints = n_constraints;
     if (n_constraints > 0)
     {
@@ -74,6 +75,7 @@ problem *new_problem(expr *objective, expr **constraints, int n_constraints)
         for (int i = 0; i < n_constraints; i++)
         {
             prob->constraints[i] = constraints[i];
+            expr_retain(constraints[i]);
         }
     }
     else
@@ -139,14 +141,11 @@ void free_problem(problem *prob)
     free(prob->gradient_values);
     free_csr_matrix(prob->stacked_jac);
 
-    /* Free expression trees with shared visited set to handle node sharing */
-    VisitedSet visited;
-    visited_init(&visited);
-
-    free_expr_tree_visited(prob->objective, &visited);
+    /* Release expression references (decrements refcount) */
+    free_expr(prob->objective);
     for (int i = 0; i < prob->n_constraints; i++)
     {
-        free_expr_tree_visited(prob->constraints[i], &visited);
+        free_expr(prob->constraints[i]);
     }
     free(prob->constraints);
 
@@ -173,6 +172,21 @@ double problem_forward(problem *prob, const double *u)
     return obj_val;
 }
 
+double *problem_constraint_forward(problem *prob, const double *u)
+{
+    /* Evaluate constraints only and copy values */
+    int offset = 0;
+    for (int i = 0; i < prob->n_constraints; i++)
+    {
+        expr *c = prob->constraints[i];
+        c->forward(c, u);
+        memcpy(prob->constraint_values + offset, c->value, c->size * sizeof(double));
+        offset += c->size;
+    }
+
+    return prob->constraint_values;
+}
+
 double *problem_gradient(problem *prob, const double *u)
 {
     /* Jacobian on objective */
diff --git a/tests/all_tests.c b/tests/all_tests.c
index 7bc8aad..751ad7c 100644
--- a/tests/all_tests.c
+++ b/tests/all_tests.c
@@ -125,6 +125,7 @@ int main(void)
     mu_run_test(test_problem_gradient, tests_run);
     mu_run_test(test_problem_jacobian, tests_run);
     mu_run_test(test_problem_jacobian_multi, tests_run);
+    mu_run_test(test_problem_constraint_forward, tests_run);
 
     printf("\n=== All %d tests passed ===\n", tests_run);
 
diff --git a/tests/problem/test_problem.h b/tests/problem/test_problem.h
index e3115d1..4ecc3d1 100644
--- a/tests/problem/test_problem.h
+++ b/tests/problem/test_problem.h
@@ -37,8 +37,12 @@ const char *test_problem_new_free(void)
     mu_assert("n_constraints wrong", prob->n_constraints == 1);
     mu_assert("total_constraint_size wrong", prob->total_constraint_size == 3);
 
-    /* Free problem (owns and frees all expressions) */
+    /* Free problem and expressions (shared ownership) */
     free_problem(prob);
+    free_expr(objective);
+    free_expr(log_x);
+    free_expr(x);
+    free_expr(x_constraint);
 
     return 0;
 }
@@ -73,6 +77,10 @@ const char *test_problem_forward(void)
     mu_assert("constraint[2] wrong", fabs(prob->constraint_values[2] - 3.0) < 1e-10);
 
     free_problem(prob);
+    free_expr(objective);
+    free_expr(log_x);
+    free_expr(x);
+    free_expr(x_constraint);
 
     return 0;
 }
@@ -99,6 +107,9 @@ const char *test_problem_gradient(void)
     mu_assert("grad[2] wrong", fabs(grad[2] - 0.25) < 1e-10);
 
     free_problem(prob);
+    free_expr(objective);
+    free_expr(log_x);
+    free_expr(x);
 
     return 0;
 }
@@ -146,6 +157,11 @@ const char *test_problem_jacobian(void)
     mu_assert("jac->x[1] wrong", fabs(jac->x[1] - 0.25) < 1e-10);
 
     free_problem(prob);
+    free_expr(objective);
+    free_expr(log_obj);
+    free_expr(x_obj);
+    free_expr(log_c1);
+    free_expr(x_c1);
 
     return 0;
 }
@@ -221,6 +237,60 @@ const char *test_problem_jacobian_multi(void)
     mu_assert("jac->x[3] wrong", fabs(jac->x[3] - exp(4.0)) < 1e-10);
 
     free_problem(prob);
+    free_expr(objective);
+    free_expr(log_obj);
+    free_expr(log_c1);
+    free_expr(exp_c2);
+    free_expr(x);
+
+    return 0;
+}
+
+/*
+ * Test problem_constraint_forward: evaluate constraints only
+ *   Constraint 1: log(x) -> [log(2), log(4)]
+ *   Constraint 2: exp(x) -> [exp(2), exp(4)]
+ */
+const char *test_problem_constraint_forward(void)
+{
+    int n_vars = 2;
+
+    /* Shared variable */
+    expr *x = new_variable(2, 1, 0, n_vars);
+
+    /* Objective: sum(log(x)) */
+    expr *log_obj = new_log(x);
+    expr *objective = new_sum(log_obj, -1);
+
+    /* Constraint 1: log(x) */
+    expr *log_c1 = new_log(x);
+
+    /* Constraint 2: exp(x) */
+    expr *exp_c2 = new_exp(x);
+
+    expr *constraints[2] = {log_c1, exp_c2};
+
+    problem *prob = new_problem(objective, constraints, 2);
+
+    double u[2] = {2.0, 4.0};
+    problem_allocate(prob, u);
+
+    double *constraint_vals = problem_constraint_forward(prob, u);
+
+    /* Check constraint values:
+     * [log(2), log(4), exp(2), exp(4)]
+     */
+    mu_assert("constraint[0] wrong", fabs(constraint_vals[0] - log(2.0)) < 1e-10);
+    mu_assert("constraint[1] wrong", fabs(constraint_vals[1] - log(4.0)) < 1e-10);
+    mu_assert("constraint[2] wrong", fabs(constraint_vals[2] - exp(2.0)) < 1e-10);
+    mu_assert("constraint[3] wrong", fabs(constraint_vals[3] - exp(4.0)) < 1e-10);
+
+    free_problem(prob);
+    free_expr(objective);
+    free_expr(log_obj);
+    free_expr(log_c1);
+    free_expr(exp_c2);
+    free_expr(x);
 
     return 0;
 }

From 4c94fb7bc3843d5ae2ba4859a4ab1f3d6f81e757 Mon Sep 17 00:00:00 2001
From: William Zijie Zhang <william@gridmatic.com>
Date: Mon, 5 Jan 2026 21:04:23 -0500
Subject: [PATCH 07/27] begin iterating on conversion of real problem, requires
 neg and promote

---
 include/affine.h                     |   2 +
 include/problem.h                    |   2 +-
 include/subexpr.h                    |   7 +
 python/bindings.c                    | 161 +++++++++--
 python/convert.py                    |  54 ++--
 python/tests/test_basic_convert.py   | 359 -------------------------
 python/tests/test_problem_convert.py | 387 ---------------------------
 src/affine/constant.c                |  15 ++
 src/affine/variable.c                |   7 +
 src/problem.c                        |  18 +-
 tests/all_tests.c                    |  10 +-
 tests/problem/test_problem.h         | 108 ++++----
 12 files changed, 274 insertions(+), 856 deletions(-)
 delete mode 100644 python/tests/test_basic_convert.py
 delete mode 100644 python/tests/test_problem_convert.py

diff --git a/include/affine.h b/include/affine.h
index e1b66ac..06a34a4 100644
--- a/include/affine.h
+++ b/include/affine.h
@@ -8,9 +8,11 @@
 expr *new_linear(expr *u, const CSR_Matrix *A);
 
 expr *new_add(expr *left, expr *right);
+expr *new_neg(expr *child);
 
 expr *new_sum(expr *child, int axis);
 expr *new_hstack(expr **args, int n_args, int n_vars);
+expr *new_promote(expr *child, int d1, int d2);
 
 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);
diff --git a/include/problem.h b/include/problem.h
index 9c29399..3dbfc9b 100644
--- a/include/problem.h
+++ b/include/problem.h
@@ -23,7 +23,7 @@ problem *new_problem(expr *objective, expr **constraints, int n_constraints);
 void problem_allocate(problem *prob, const double *u);
 void free_problem(problem *prob);
 
-double problem_forward(problem *prob, const double *u);
+double problem_objective_forward(problem *prob, const double *u);
 double *problem_constraint_forward(problem *prob, const double *u);
 double *problem_gradient(problem *prob, const double *u);
 CSR_Matrix *problem_jacobian(problem *prob, const double *u);
diff --git a/include/subexpr.h b/include/subexpr.h
index 033fd7f..3201877 100644
--- a/include/subexpr.h
+++ b/include/subexpr.h
@@ -49,4 +49,11 @@ typedef struct hstack_expr
     CSR_Matrix *CSR_work; /* for summing Hessians of children */
 } hstack_expr;
 
+/* Promote (broadcast) to larger shape */
+typedef struct promote_expr
+{
+    expr base;
+    /* target shape stored in base.d1, base.d2 */
+} promote_expr;
+
 #endif /* SUBEXPR_H */
diff --git a/python/bindings.c b/python/bindings.c
index 06354a3..3f3ef31 100644
--- a/python/bindings.c
+++ b/python/bindings.c
@@ -48,6 +48,88 @@ static PyObject *py_make_variable(PyObject *self, PyObject *args)
     return PyCapsule_New(node, EXPR_CAPSULE_NAME, expr_capsule_destructor);
 }
 
+static PyObject *py_make_constant(PyObject *self, PyObject *args)
+{
+    int d1, d2, n_vars;
+    PyObject *values_obj;
+    if (!PyArg_ParseTuple(args, "iiiO", &d1, &d2, &n_vars, &values_obj))
+    {
+        return NULL;
+    }
+
+    PyArrayObject *values_array =
+        (PyArrayObject *) PyArray_FROM_OTF(values_obj, NPY_DOUBLE, NPY_ARRAY_IN_ARRAY);
+    if (!values_array)
+    {
+        return NULL;
+    }
+
+    expr *node =
+        new_constant(d1, d2, n_vars, (const double *) PyArray_DATA(values_array));
+    Py_DECREF(values_array);
+
+    if (!node)
+    {
+        PyErr_SetString(PyExc_RuntimeError, "failed to create constant node");
+        return NULL;
+    }
+    return PyCapsule_New(node, EXPR_CAPSULE_NAME, expr_capsule_destructor);
+}
+
+static PyObject *py_make_linear(PyObject *self, PyObject *args)
+{
+    PyObject *child_capsule;
+    PyObject *data_obj, *indices_obj, *indptr_obj;
+    int m, n;
+    if (!PyArg_ParseTuple(args, "OOOOii", &child_capsule, &data_obj, &indices_obj,
+                          &indptr_obj, &m, &n))
+    {
+        return NULL;
+    }
+
+    expr *child = (expr *) PyCapsule_GetPointer(child_capsule, EXPR_CAPSULE_NAME);
+    if (!child)
+    {
+        PyErr_SetString(PyExc_ValueError, "invalid child capsule");
+        return NULL;
+    }
+
+    PyArrayObject *data_array =
+        (PyArrayObject *) PyArray_FROM_OTF(data_obj, NPY_DOUBLE, NPY_ARRAY_IN_ARRAY);
+    PyArrayObject *indices_array =
+        (PyArrayObject *) PyArray_FROM_OTF(indices_obj, NPY_INT32, NPY_ARRAY_IN_ARRAY);
+    PyArrayObject *indptr_array =
+        (PyArrayObject *) PyArray_FROM_OTF(indptr_obj, NPY_INT32, NPY_ARRAY_IN_ARRAY);
+
+    if (!data_array || !indices_array || !indptr_array)
+    {
+        Py_XDECREF(data_array);
+        Py_XDECREF(indices_array);
+        Py_XDECREF(indptr_array);
+        return NULL;
+    }
+
+    int nnz = (int) PyArray_SIZE(data_array);
+    CSR_Matrix *A = new_csr_matrix(m, n, nnz);
+    memcpy(A->x, PyArray_DATA(data_array), nnz * sizeof(double));
+    memcpy(A->i, PyArray_DATA(indices_array), nnz * sizeof(int));
+    memcpy(A->p, PyArray_DATA(indptr_array), (m + 1) * sizeof(int));
+
+    Py_DECREF(data_array);
+    Py_DECREF(indices_array);
+    Py_DECREF(indptr_array);
+
+    expr *node = new_linear(child, A);
+    free_csr_matrix(A);
+
+    if (!node)
+    {
+        PyErr_SetString(PyExc_RuntimeError, "failed to create linear node");
+        return NULL;
+    }
+    return PyCapsule_New(node, EXPR_CAPSULE_NAME, expr_capsule_destructor);
+}
+
 static PyObject *py_make_log(PyObject *self, PyObject *args)
 {
     PyObject *child_capsule;
@@ -142,6 +224,53 @@ static PyObject *py_make_sum(PyObject *self, PyObject *args)
     return PyCapsule_New(node, EXPR_CAPSULE_NAME, expr_capsule_destructor);
 }
 
+static PyObject *py_make_neg(PyObject *self, PyObject *args)
+{
+    PyObject *child_capsule;
+    if (!PyArg_ParseTuple(args, "O", &child_capsule))
+    {
+        return NULL;
+    }
+    expr *child = (expr *) PyCapsule_GetPointer(child_capsule, EXPR_CAPSULE_NAME);
+    if (!child)
+    {
+        PyErr_SetString(PyExc_ValueError, "invalid child capsule");
+        return NULL;
+    }
+
+    expr *node = new_neg(child);
+    if (!node)
+    {
+        PyErr_SetString(PyExc_RuntimeError, "failed to create neg node");
+        return NULL;
+    }
+    return PyCapsule_New(node, EXPR_CAPSULE_NAME, expr_capsule_destructor);
+}
+
+static PyObject *py_make_promote(PyObject *self, PyObject *args)
+{
+    PyObject *child_capsule;
+    int d1, d2;
+    if (!PyArg_ParseTuple(args, "Oii", &child_capsule, &d1, &d2))
+    {
+        return NULL;
+    }
+    expr *child = (expr *) PyCapsule_GetPointer(child_capsule, EXPR_CAPSULE_NAME);
+    if (!child)
+    {
+        PyErr_SetString(PyExc_ValueError, "invalid child capsule");
+        return NULL;
+    }
+
+    expr *node = new_promote(child, d1, d2);
+    if (!node)
+    {
+        PyErr_SetString(PyExc_RuntimeError, "failed to create promote node");
+        return NULL;
+    }
+    return PyCapsule_New(node, EXPR_CAPSULE_NAME, expr_capsule_destructor);
+}
+
 static PyObject *py_forward(PyObject *self, PyObject *args)
 {
     PyObject *node_capsule;
@@ -338,7 +467,7 @@ static PyObject *py_problem_allocate(PyObject *self, PyObject *args)
     Py_RETURN_NONE;
 }
 
-static PyObject *py_problem_forward(PyObject *self, PyObject *args)
+static PyObject *py_problem_objective_forward(PyObject *self, PyObject *args)
 {
     PyObject *prob_capsule;
     PyObject *u_obj;
@@ -362,30 +491,10 @@ static PyObject *py_problem_forward(PyObject *self, PyObject *args)
         return NULL;
     }
 
-    double obj_val = problem_forward(prob, (const double *) PyArray_DATA(u_array));
-
-    // Create constraint values array
-    PyObject *constraint_vals = NULL;
-    if (prob->total_constraint_size > 0)
-    {
-        npy_intp size = prob->total_constraint_size;
-        constraint_vals = PyArray_SimpleNew(1, &size, NPY_DOUBLE);
-        if (!constraint_vals)
-        {
-            Py_DECREF(u_array);
-            return NULL;
-        }
-        memcpy(PyArray_DATA((PyArrayObject *) constraint_vals), prob->constraint_values,
-               size * sizeof(double));
-    }
-    else
-    {
-        npy_intp size = 0;
-        constraint_vals = PyArray_SimpleNew(1, &size, NPY_DOUBLE);
-    }
+    double obj_val = problem_objective_forward(prob, (const double *) PyArray_DATA(u_array));
 
     Py_DECREF(u_array);
-    return Py_BuildValue("(dO)", obj_val, constraint_vals);
+    return Py_BuildValue("d", obj_val);
 }
 
 static PyObject *py_problem_constraint_forward(PyObject *self, PyObject *args)
@@ -541,15 +650,19 @@ static PyObject *py_problem_jacobian(PyObject *self, PyObject *args)
 
 static PyMethodDef DNLPMethods[] = {
     {"make_variable", py_make_variable, METH_VARARGS, "Create variable node"},
+    {"make_constant", py_make_constant, METH_VARARGS, "Create constant node"},
+    {"make_linear", py_make_linear, METH_VARARGS, "Create linear op node"},
     {"make_log", py_make_log, METH_VARARGS, "Create log node"},
     {"make_exp", py_make_exp, METH_VARARGS, "Create exp node"},
     {"make_add", py_make_add, METH_VARARGS, "Create add node"},
     {"make_sum", py_make_sum, METH_VARARGS, "Create sum node"},
+    {"make_neg", py_make_neg, METH_VARARGS, "Create neg node"},
+    {"make_promote", py_make_promote, METH_VARARGS, "Create promote node"},
     {"forward", py_forward, METH_VARARGS, "Run forward pass and return values"},
     {"jacobian", py_jacobian, METH_VARARGS, "Compute jacobian and return CSR components"},
     {"make_problem", py_make_problem, METH_VARARGS, "Create problem from objective and constraints"},
     {"problem_allocate", py_problem_allocate, METH_VARARGS, "Allocate problem resources"},
-    {"problem_forward", py_problem_forward, METH_VARARGS, "Evaluate objective and constraints"},
+    {"problem_objective_forward", py_problem_objective_forward, METH_VARARGS, "Evaluate objective only"},
     {"problem_constraint_forward", py_problem_constraint_forward, METH_VARARGS, "Evaluate constraints only"},
     {"problem_gradient", py_problem_gradient, METH_VARARGS, "Compute objective gradient"},
     {"problem_jacobian", py_problem_jacobian, METH_VARARGS, "Compute constraint jacobian"},
diff --git a/python/convert.py b/python/convert.py
index 900a32d..a7a6166 100644
--- a/python/convert.py
+++ b/python/convert.py
@@ -53,19 +53,39 @@ def build_variable_dict(variables: list) -> tuple[dict, int]:
         c_var = diffengine.make_variable(d1, d2, offset, n_vars)
         var_dict[var.id] = c_var
 
-    return var_dict
+    return var_dict, n_vars
 
 
-def _convert_expr(expr, var_dict: dict):
+def _convert_expr(expr, var_dict: dict, n_vars: int):
     """Convert CVXPY expression using pre-built variable dictionary."""
     # Base case: variable lookup
     if isinstance(expr, cp.Variable):
         return var_dict[expr.id]
 
+    # Base case: constant
+    if isinstance(expr, cp.Constant):
+        value = np.asarray(expr.value, dtype=np.float64).flatten()
+        d1 = expr.shape[0] if len(expr.shape) >= 1 else 1
+        d2 = expr.shape[1] if len(expr.shape) >= 2 else 1
+        return diffengine.make_constant(d1, d2, n_vars, value)
+
     # Recursive case: atoms
     atom_name = type(expr).__name__
+
+    # Handle NegExpression using neg atom
+    if atom_name == "NegExpression":
+        child = _convert_expr(expr.args[0], var_dict, n_vars)
+        return diffengine.make_neg(child)
+
+    # Handle Promote (broadcasts scalar/vector to larger shape)
+    if atom_name == "Promote":
+        child = _convert_expr(expr.args[0], var_dict, n_vars)
+        d1 = expr.shape[0] if len(expr.shape) >= 1 else 1
+        d2 = expr.shape[1] if len(expr.shape) >= 2 else 1
+        return diffengine.make_promote(child, d1, d2)
+
     if atom_name in ATOM_CONVERTERS:
-        children = [_convert_expr(arg, var_dict) for arg in expr.args]
+        children = [_convert_expr(arg, var_dict, n_vars) for arg in expr.args]
         converter = ATOM_CONVERTERS[atom_name]
         # N-ary ops (like AddExpression) take list, unary ops take single arg
         if atom_name == "AddExpression":
@@ -86,21 +106,21 @@ def convert_expressions(problem: cp.Problem) -> tuple:
         c_objective: C expression for objective
         c_constraints: list of C expressions for constraints
     """
-    var_dict = build_variable_dict(problem.variables())
+    var_dict, n_vars = build_variable_dict(problem.variables())
 
     # Convert objective
-    c_objective = _convert_expr(problem.objective.expr, var_dict)
+    c_objective = _convert_expr(problem.objective.expr, var_dict, n_vars)
 
     # Convert constraints (expression part only for now)
     c_constraints = []
     for constr in problem.constraints:
-        c_expr = _convert_expr(constr.expr, var_dict)
+        c_expr = _convert_expr(constr.expr, var_dict, n_vars)
         c_constraints.append(c_expr)
 
     return c_objective, c_constraints
 
 
-def convert_problem(problem: cp.Problem) -> "Problem":
+def convert_problem(problem: cp.Problem) -> "C_problem":
     """
     Convert CVXPY Problem to C problem struct.
 
@@ -108,18 +128,20 @@ def convert_problem(problem: cp.Problem) -> "Problem":
         problem: CVXPY Problem object
 
     Returns:
-        Problem wrapper around C problem struct
+        C_Problem wrapper around C problem struct
     """
-    return Problem(problem)
+    return C_problem(problem)
 
 
-class Problem:
+class C_problem:
     """Wrapper around C problem struct for CVXPY problems."""
 
     def __init__(self, cvxpy_problem: cp.Problem):
-        var_dict = build_variable_dict(cvxpy_problem.variables())
-        c_obj = _convert_expr(cvxpy_problem.objective.expr, var_dict)
-        c_constraints = [_convert_expr(c.expr, var_dict) for c in cvxpy_problem.constraints]
+        var_dict, n_vars = build_variable_dict(cvxpy_problem.variables())
+        c_obj = _convert_expr(cvxpy_problem.objective.expr, var_dict, n_vars)
+        c_constraints = [
+            _convert_expr(c.expr, var_dict, n_vars) for c in cvxpy_problem.constraints
+        ]
         self._capsule = diffengine.make_problem(c_obj, c_constraints)
         self._allocated = False
 
@@ -128,9 +150,9 @@ def allocate(self, u: np.ndarray):
         diffengine.problem_allocate(self._capsule, u)
         self._allocated = True
 
-    def forward(self, u: np.ndarray) -> tuple[float, np.ndarray]:
-        """Evaluate objective and constraints. Returns (obj_value, constraint_values)."""
-        return diffengine.problem_forward(self._capsule, u)
+    def objective_forward(self, u: np.ndarray) -> float:
+        """Evaluate objective. Returns obj_value float."""
+        return diffengine.problem_objective_forward(self._capsule, u)
 
     def constraint_forward(self, u: np.ndarray) -> np.ndarray:
         """Evaluate constraints only. Returns constraint_values array."""
diff --git a/python/tests/test_basic_convert.py b/python/tests/test_basic_convert.py
deleted file mode 100644
index 40318d8..0000000
--- a/python/tests/test_basic_convert.py
+++ /dev/null
@@ -1,359 +0,0 @@
-import cvxpy as cp
-import numpy as np
-from scipy import sparse
-import DNLP_diff_engine as diffengine
-from convert import convert_expressions, get_jacobian
-
-
-def test_sum_log():
-    """Test sum(log(x)) forward and jacobian."""
-    x = cp.Variable(4)
-    problem = cp.Problem(cp.Minimize(cp.sum(cp.log(x))))
-    c_obj, _ = convert_expressions(problem)
-
-    test_values = np.array([1.0, 2.0, 3.0, 4.0])
-
-    # Forward
-    result = diffengine.forward(c_obj, test_values)
-    expected = np.sum(np.log(test_values))
-    assert np.allclose(result, expected)
-
-    # Jacobian: d/dx sum(log(x)) = [1/x_1, 1/x_2, ...]
-    jac = get_jacobian(c_obj, test_values)
-    expected_jac = (1.0 / test_values).reshape(1, -1)
-    assert np.allclose(jac.toarray(), expected_jac)
-
-
-def test_sum_exp():
-    """Test sum(exp(x)) forward and jacobian."""
-    x = cp.Variable(3)
-    problem = cp.Problem(cp.Minimize(cp.sum(cp.exp(x))))
-    c_obj, _ = convert_expressions(problem)
-
-    test_values = np.array([0.0, 1.0, 2.0])
-
-    # Forward
-    result = diffengine.forward(c_obj, test_values)
-    expected = np.sum(np.exp(test_values))
-    assert np.allclose(result, expected)
-
-    # Jacobian: d/dx sum(exp(x)) = [exp(x_1), exp(x_2), ...]
-    jac = get_jacobian(c_obj, test_values)
-    expected_jac = np.exp(test_values).reshape(1, -1)
-    assert np.allclose(jac.toarray(), expected_jac)
-
-
-def test_two_variables_elementwise_add():
-    """Test sum(log(x + y)) - elementwise after add."""
-    x = cp.Variable(2)
-    y = cp.Variable(2)
-    problem = cp.Problem(cp.Minimize(cp.sum(cp.log(x + y))))
-    c_obj, _ = convert_expressions(problem)
-
-    test_values = np.array([1.0, 2.0, 3.0, 4.0])
-
-    # Forward
-    result = diffengine.forward(c_obj, test_values)
-    expected = np.sum(np.log(np.array([1+3, 2+4])))
-    assert np.allclose(result, expected)
-
-    # TODO: Jacobian for elementwise(add(...)) patterns not yet supported
-
-
-def test_variable_reuse():
-    """Test sum(log(x) + exp(x)) - same variable used twice."""
-    x = cp.Variable(2)
-    problem = cp.Problem(cp.Minimize(cp.sum(cp.log(x) + cp.exp(x))))
-    c_obj, _ = convert_expressions(problem)
-
-    test_values = np.array([1.0, 2.0])
-
-    # Forward
-    result = diffengine.forward(c_obj, test_values)
-    expected = np.sum(np.log(test_values) + np.exp(test_values))
-    assert np.allclose(result, expected)
-
-    # Jacobian: d/dx_i = 1/x_i + exp(x_i)
-    jac = get_jacobian(c_obj, test_values)
-    expected_jac = (1.0 / test_values + np.exp(test_values)).reshape(1, -1)
-    assert np.allclose(jac.toarray(), expected_jac)
-
-
-def test_four_variables_elementwise_add():
-    """Test sum(log(a + b) + exp(c + d)) - elementwise after add."""
-    a = cp.Variable(3)
-    b = cp.Variable(3)
-    c = cp.Variable(3)
-    d = cp.Variable(3)
-    problem = cp.Problem(cp.Minimize(cp.sum(cp.log(a + b) + cp.exp(c + d))))
-    c_obj, _ = convert_expressions(problem)
-
-    a_vals = np.array([1.0, 2.0, 3.0])
-    b_vals = np.array([0.5, 1.0, 1.5])
-    c_vals = np.array([0.1, 0.2, 0.3])
-    d_vals = np.array([0.1, 0.1, 0.1])
-    test_values = np.concatenate([a_vals, b_vals, c_vals, d_vals])
-
-    # Forward
-    result = diffengine.forward(c_obj, test_values)
-    expected = np.sum(np.log(a_vals + b_vals) + np.exp(c_vals + d_vals))
-    assert np.allclose(result, expected)
-
-    # TODO: Jacobian for elementwise(add(...)) patterns not yet supported
-
-
-def test_deep_nesting():
-    """Test sum(log(exp(log(exp(x))))) - deeply nested elementwise."""
-    x = cp.Variable(4)
-    problem = cp.Problem(cp.Minimize(cp.sum(cp.log(cp.exp(cp.log(cp.exp(x)))))))
-    c_obj, _ = convert_expressions(problem)
-
-    test_values = np.array([0.5, 1.0, 1.5, 2.0])
-
-    # Forward
-    result = diffengine.forward(c_obj, test_values)
-    expected = np.sum(np.log(np.exp(np.log(np.exp(test_values)))))
-    assert np.allclose(result, expected)
-
-    # TODO: Jacobian for nested elementwise compositions not yet supported
-
-
-def test_chained_additions():
-    """Test sum(x + y + z + w) - chained additions."""
-    x = cp.Variable(2)
-    y = cp.Variable(2)
-    z = cp.Variable(2)
-    w = cp.Variable(2)
-    problem = cp.Problem(cp.Minimize(cp.sum(x + y + z + w)))
-    c_obj, _ = convert_expressions(problem)
-
-    x_vals = np.array([1.0, 2.0])
-    y_vals = np.array([3.0, 4.0])
-    z_vals = np.array([5.0, 6.0])
-    w_vals = np.array([7.0, 8.0])
-    test_values = np.concatenate([x_vals, y_vals, z_vals, w_vals])
-
-    # Forward
-    result = diffengine.forward(c_obj, test_values)
-    expected = np.sum(x_vals + y_vals + z_vals + w_vals)
-    assert np.allclose(result, expected)
-
-    # TODO: Jacobian for sum(add(...)) patterns not yet supported
-
-
-def test_variable_used_multiple_times():
-    """Test sum(log(x) + exp(x) + x) - variable used 3+ times."""
-    x = cp.Variable(3)
-    problem = cp.Problem(cp.Minimize(cp.sum(cp.log(x) + cp.exp(x) + x)))
-    c_obj, _ = convert_expressions(problem)
-
-    test_values = np.array([1.0, 2.0, 3.0])
-
-    # Forward
-    result = diffengine.forward(c_obj, test_values)
-    expected = np.sum(np.log(test_values) + np.exp(test_values) + test_values)
-    assert np.allclose(result, expected)
-
-    # TODO: Jacobian for expressions with sum(variable) not yet supported
-
-
-def test_larger_variable():
-    """Test sum(log(x)) with larger variable (100 elements)."""
-    x = cp.Variable(100)
-    problem = cp.Problem(cp.Minimize(cp.sum(cp.log(x))))
-    c_obj, _ = convert_expressions(problem)
-
-    test_values = np.linspace(1.0, 10.0, 100)
-
-    # Forward
-    result = diffengine.forward(c_obj, test_values)
-    expected = np.sum(np.log(test_values))
-    assert np.allclose(result, expected)
-
-    # Jacobian
-    jac = get_jacobian(c_obj, test_values)
-    expected_jac = (1.0 / test_values).reshape(1, -1)
-    assert np.allclose(jac.toarray(), expected_jac)
-
-
-def test_matrix_variable():
-    """Test sum(log(X)) with 2D matrix variable (3x4)."""
-    X = cp.Variable((3, 4))
-    problem = cp.Problem(cp.Minimize(cp.sum(cp.log(X))))
-    c_obj, _ = convert_expressions(problem)
-
-    test_values = np.arange(1.0, 13.0)  # 12 elements
-
-    # Forward
-    result = diffengine.forward(c_obj, test_values)
-    expected = np.sum(np.log(test_values))
-    assert np.allclose(result, expected)
-
-    # Jacobian
-    jac = get_jacobian(c_obj, test_values)
-    expected_jac = (1.0 / test_values).reshape(1, -1)
-    assert np.allclose(jac.toarray(), expected_jac)
-
-
-def test_mixed_sizes():
-    """Test sum(log(a)) + sum(log(b)) + sum(log(c)) with different sized variables."""
-    a = cp.Variable(2)
-    b = cp.Variable(5)
-    c = cp.Variable(3)
-    problem = cp.Problem(cp.Minimize(cp.sum(cp.log(a)) + cp.sum(cp.log(b)) + cp.sum(cp.log(c))))
-    c_obj, _ = convert_expressions(problem)
-
-    a_vals = np.array([1.0, 2.0])
-    b_vals = np.array([1.0, 2.0, 3.0, 4.0, 5.0])
-    c_vals = np.array([1.0, 2.0, 3.0])
-    test_values = np.concatenate([a_vals, b_vals, c_vals])
-
-    # Forward
-    result = diffengine.forward(c_obj, test_values)
-    expected = np.sum(np.log(a_vals)) + np.sum(np.log(b_vals)) + np.sum(np.log(c_vals))
-    assert np.allclose(result, expected)
-
-    # Jacobian
-    jac = get_jacobian(c_obj, test_values)
-    expected_jac = (1.0 / test_values).reshape(1, -1)
-    assert np.allclose(jac.toarray(), expected_jac)
-
-
-def test_multiple_variables_log_exp():
-    """Test sum(log(a)) + sum(log(b)) + sum(exp(c)) + sum(exp(d))."""
-    a = cp.Variable(2)
-    b = cp.Variable(2)
-    c = cp.Variable(2)
-    d = cp.Variable(2)
-    obj = cp.sum(cp.log(a)) + cp.sum(cp.log(b)) + cp.sum(cp.exp(c)) + cp.sum(cp.exp(d))
-    problem = cp.Problem(cp.Minimize(obj))
-    c_obj, _ = convert_expressions(problem)
-
-    a_vals = np.array([1.0, 2.0])
-    b_vals = np.array([0.5, 1.0])
-    c_vals = np.array([0.1, 0.2])
-    d_vals = np.array([0.1, 0.1])
-    test_values = np.concatenate([a_vals, b_vals, c_vals, d_vals])
-
-    # Forward
-    result = diffengine.forward(c_obj, test_values)
-    expected = (np.sum(np.log(a_vals)) + np.sum(np.log(b_vals)) +
-                np.sum(np.exp(c_vals)) + np.sum(np.exp(d_vals)))
-    assert np.allclose(result, expected)
-
-    # Jacobian
-    jac = get_jacobian(c_obj, test_values)
-    df_da = 1.0 / a_vals
-    df_db = 1.0 / b_vals
-    df_dc = np.exp(c_vals)
-    df_dd = np.exp(d_vals)
-    expected_jac = np.concatenate([df_da, df_db, df_dc, df_dd]).reshape(1, -1)
-    assert np.allclose(jac.toarray(), expected_jac)
-
-
-def test_two_variables_separate_sums():
-    """Test sum(log(x) + log(y)) - two variables with separate elementwise ops."""
-    x = cp.Variable(2)
-    y = cp.Variable(2)
-    problem = cp.Problem(cp.Minimize(cp.sum(cp.log(x) + cp.log(y))))
-    c_obj, _ = convert_expressions(problem)
-
-    test_values = np.array([1.0, 2.0, 3.0, 4.0])
-
-    # Forward
-    result = diffengine.forward(c_obj, test_values)
-    expected = np.sum(np.log(test_values[:2]) + np.log(test_values[2:]))
-    assert np.allclose(result, expected)
-
-    # Jacobian
-    jac = get_jacobian(c_obj, test_values)
-    expected_jac = np.array([[1/1, 1/2, 1/3, 1/4]])
-    assert np.allclose(jac.toarray(), expected_jac)
-
-
-def test_complex_objective_elementwise_add():
-    """Test sum(log(x + y)) + sum(exp(y + z)) + sum(log(z + x)) - elementwise after add."""
-    x = cp.Variable(3)
-    y = cp.Variable(3)
-    z = cp.Variable(3)
-    obj = cp.sum(cp.log(x + y)) + cp.sum(cp.exp(y + z)) + cp.sum(cp.log(z + x))
-    problem = cp.Problem(cp.Minimize(obj))
-    c_obj, _ = convert_expressions(problem)
-
-    x_vals = np.array([1.0, 2.0, 3.0])
-    y_vals = np.array([0.5, 1.0, 1.5])
-    z_vals = np.array([0.2, 0.3, 0.4])
-    test_values = np.concatenate([x_vals, y_vals, z_vals])
-
-    # Forward
-    result = diffengine.forward(c_obj, test_values)
-    expected = (np.sum(np.log(x_vals + y_vals)) +
-                np.sum(np.exp(y_vals + z_vals)) +
-                np.sum(np.log(z_vals + x_vals)))
-    assert np.allclose(result, expected)
-
-    # TODO: Jacobian for elementwise(add(...)) patterns not yet supported
-
-
-def test_complex_objective_no_add():
-    """Test sum(log(x) + exp(y) + log(z)) - multiple elementwise ops without add composition."""
-    x = cp.Variable(2)
-    y = cp.Variable(2)
-    z = cp.Variable(2)
-    obj = cp.sum(cp.log(x) + cp.exp(y) + cp.log(z))
-    problem = cp.Problem(cp.Minimize(obj))
-    c_obj, _ = convert_expressions(problem)
-
-    x_vals = np.array([1.0, 2.0])
-    y_vals = np.array([0.5, 1.0])
-    z_vals = np.array([0.2, 0.3])
-    test_values = np.concatenate([x_vals, y_vals, z_vals])
-
-    # Forward
-    result = diffengine.forward(c_obj, test_values)
-    expected = np.sum(np.log(x_vals) + np.exp(y_vals) + np.log(z_vals))
-    assert np.allclose(result, expected)
-
-    # Jacobian
-    jac = get_jacobian(c_obj, test_values)
-    df_dx = 1.0 / x_vals
-    df_dy = np.exp(y_vals)
-    df_dz = 1.0 / z_vals
-    expected_jac = np.concatenate([df_dx, df_dy, df_dz]).reshape(1, -1)
-    assert np.allclose(jac.toarray(), expected_jac)
-
-
-def test_log_exp_identity():
-    """Test sum(log(exp(x))) = sum(x) identity - nested elementwise."""
-    x = cp.Variable(5)
-    problem = cp.Problem(cp.Minimize(cp.sum(cp.log(cp.exp(x)))))
-    c_obj, _ = convert_expressions(problem)
-
-    test_values = np.array([-1.0, 0.0, 1.0, 2.0, 3.0])
-
-    # Forward
-    result = diffengine.forward(c_obj, test_values)
-    expected = np.sum(test_values)  # log(exp(x)) = x
-    assert np.allclose(result, expected)
-
-    # TODO: Jacobian for nested elementwise compositions not yet supported
-
-
-if __name__ == "__main__":
-    test_sum_log()
-    test_sum_exp()
-    test_two_variables_elementwise_add()
-    test_variable_reuse()
-    test_four_variables_elementwise_add()
-    test_deep_nesting()
-    test_chained_additions()
-    test_variable_used_multiple_times()
-    test_larger_variable()
-    test_matrix_variable()
-    test_mixed_sizes()
-    test_multiple_variables_log_exp()
-    test_two_variables_separate_sums()
-    test_complex_objective_elementwise_add()
-    test_complex_objective_no_add()
-    test_log_exp_identity()
-    print("All tests passed!")
diff --git a/python/tests/test_problem_convert.py b/python/tests/test_problem_convert.py
deleted file mode 100644
index de605cd..0000000
--- a/python/tests/test_problem_convert.py
+++ /dev/null
@@ -1,387 +0,0 @@
-import cvxpy as cp
-import numpy as np
-from scipy import sparse
-import DNLP_diff_engine as diffengine
-from convert import Problem
-
-
-# ============ Low-level problem struct tests ============
-
-def test_problem_forward_lowlevel():
-    """Test problem_forward for objective and constraint values (low-level)."""
-    n_vars = 3
-    x_obj = diffengine.make_variable(n_vars, 1, 0, n_vars)
-    log_obj = diffengine.make_log(x_obj)
-    objective = diffengine.make_sum(log_obj, -1)
-
-    x_c = diffengine.make_variable(n_vars, 1, 0, n_vars)
-    log_c = diffengine.make_log(x_c)
-    constraints = [log_c]
-
-    prob = diffengine.make_problem(objective, constraints)
-    u = np.array([1.0, 2.0, 3.0])
-    diffengine.problem_allocate(prob, u)
-
-    obj_val, constraint_vals = diffengine.problem_forward(prob, u)
-
-    expected_obj = np.sum(np.log(u))
-    assert np.allclose(obj_val, expected_obj)
-    assert np.allclose(constraint_vals, np.log(u))
-
-
-def test_problem_gradient_lowlevel():
-    """Test problem_gradient for objective gradient (low-level)."""
-    n_vars = 3
-    x = diffengine.make_variable(n_vars, 1, 0, n_vars)
-    log_x = diffengine.make_log(x)
-    objective = diffengine.make_sum(log_x, -1)
-
-    prob = diffengine.make_problem(objective, [])
-    u = np.array([1.0, 2.0, 4.0])
-    diffengine.problem_allocate(prob, u)
-
-    grad = diffengine.problem_gradient(prob, u)
-    expected_grad = 1.0 / u
-    assert np.allclose(grad, expected_grad)
-
-
-def test_problem_jacobian_lowlevel():
-    """Test problem_jacobian for constraint jacobian (low-level)."""
-    n_vars = 2
-    x_obj = diffengine.make_variable(n_vars, 1, 0, n_vars)
-    log_obj = diffengine.make_log(x_obj)
-    objective = diffengine.make_sum(log_obj, -1)
-
-    x_c = diffengine.make_variable(n_vars, 1, 0, n_vars)
-    log_c = diffengine.make_log(x_c)
-    constraints = [log_c]
-
-    prob = diffengine.make_problem(objective, constraints)
-    u = np.array([2.0, 4.0])
-    diffengine.problem_allocate(prob, u)
-
-    data, indices, indptr, shape = diffengine.problem_jacobian(prob, u)
-    jac = sparse.csr_matrix((data, indices, indptr), shape=shape)
-
-    expected_jac = np.diag(1.0 / u)
-    assert np.allclose(jac.toarray(), expected_jac)
-
-
-def test_problem_constraint_forward_lowlevel():
-    """Test problem_constraint_forward for constraint values only (low-level)."""
-    n_vars = 2
-    x = diffengine.make_variable(n_vars, 1, 0, n_vars)
-
-    log_obj = diffengine.make_log(x)
-    objective = diffengine.make_sum(log_obj, -1)
-
-    log_c = diffengine.make_log(x)
-    exp_c = diffengine.make_exp(x)
-    constraints = [log_c, exp_c]
-
-    prob = diffengine.make_problem(objective, constraints)
-    u = np.array([2.0, 4.0])
-    diffengine.problem_allocate(prob, u)
-
-    constraint_vals = diffengine.problem_constraint_forward(prob, u)
-
-    # Expected: [log(2), log(4), exp(2), exp(4)]
-    expected = np.concatenate([np.log(u), np.exp(u)])
-    assert np.allclose(constraint_vals, expected)
-
-
-def test_problem_no_constraints_lowlevel():
-    """Test Problem with no constraints (low-level)."""
-    n_vars = 3
-    x = diffengine.make_variable(n_vars, 1, 0, n_vars)
-    log_x = diffengine.make_log(x)
-    objective = diffengine.make_sum(log_x, -1)
-
-    prob = diffengine.make_problem(objective, [])
-    u = np.array([1.0, 2.0, 3.0])
-    diffengine.problem_allocate(prob, u)
-
-    obj_val, constraint_vals = diffengine.problem_forward(prob, u)
-    assert np.allclose(obj_val, np.sum(np.log(u)))
-    assert len(constraint_vals) == 0
-
-    grad = diffengine.problem_gradient(prob, u)
-    assert np.allclose(grad, 1.0 / u)
-
-    data, indices, indptr, shape = diffengine.problem_jacobian(prob, u)
-    jac = sparse.csr_matrix((data, indices, indptr), shape=shape)
-    assert jac.shape == (0, 3)
-
-
-# ============ Problem class tests using convert ============
-
-def test_problem_single_constraint():
-    """Test Problem class with single constraint using convert."""
-    x = cp.Variable(3)
-    obj = cp.sum(cp.log(x))
-    constraints = [cp.log(x)]  # log(x) as constraint expression
-
-    cvxpy_prob = cp.Problem(cp.Minimize(obj), constraints)
-    prob = Problem(cvxpy_prob)
-
-    u = np.array([1.0, 2.0, 3.0])
-    prob.allocate(u)
-
-    # Test forward
-    obj_val, constraint_vals = prob.forward(u)
-    assert np.allclose(obj_val, np.sum(np.log(u)))
-    assert np.allclose(constraint_vals, np.log(u))
-
-    # Test gradient
-    grad = prob.gradient(u)
-    assert np.allclose(grad, 1.0 / u)
-
-    # Test jacobian
-    jac = prob.jacobian(u)
-    expected_jac = np.diag(1.0 / u)
-    assert np.allclose(jac.toarray(), expected_jac)
-
-
-def test_problem_two_constraints():
-    """Test Problem class with two constraints."""
-    x = cp.Variable(2)
-    obj = cp.sum(cp.log(x))
-    constraints = [
-        cp.log(x),  # constraint 1: log(x), size 2
-        cp.exp(x),  # constraint 2: exp(x), size 2
-    ]
-
-    cvxpy_prob = cp.Problem(cp.Minimize(obj), constraints)
-    prob = Problem(cvxpy_prob)
-
-    u = np.array([1.0, 2.0])
-    prob.allocate(u)
-
-    # Test forward
-    obj_val, constraint_vals = prob.forward(u)
-    assert np.allclose(obj_val, np.sum(np.log(u)))
-    # Constraint values should be stacked: [log(1), log(2), exp(1), exp(2)]
-    expected_constraint_vals = np.concatenate([np.log(u), np.exp(u)])
-    assert np.allclose(constraint_vals, expected_constraint_vals)
-
-    # Test gradient
-    grad = prob.gradient(u)
-    assert np.allclose(grad, 1.0 / u)
-
-    # Test jacobian - stacked vertically
-    jac = prob.jacobian(u)
-    assert jac.shape == (4, 2)  # 2 constraints * 2 elements each
-    # First 2 rows: d(log(x))/dx = diag(1/x)
-    # Last 2 rows: d(exp(x))/dx = diag(exp(x))
-    expected_jac = np.vstack([np.diag(1.0 / u), np.diag(np.exp(u))])
-    assert np.allclose(jac.toarray(), expected_jac)
-
-
-def test_problem_three_constraints_different_sizes():
-    """Test Problem with three constraints of different types."""
-    x = cp.Variable(3)
-    obj = cp.sum(cp.exp(x))
-    constraints = [
-        cp.log(x),       # size 3
-        cp.exp(x),       # size 3
-        cp.sum(cp.log(x)),  # size 1 (scalar)
-    ]
-
-    cvxpy_prob = cp.Problem(cp.Minimize(obj), constraints)
-    prob = Problem(cvxpy_prob)
-
-    u = np.array([1.0, 2.0, 3.0])
-    prob.allocate(u)
-
-    # Test forward
-    obj_val, constraint_vals = prob.forward(u)
-    assert np.allclose(obj_val, np.sum(np.exp(u)))
-    # Constraint values: [log(x), exp(x), sum(log(x))]
-    expected_constraint_vals = np.concatenate([
-        np.log(u),
-        np.exp(u),
-        [np.sum(np.log(u))]
-    ])
-    assert np.allclose(constraint_vals, expected_constraint_vals)
-
-    # Test gradient of sum(exp(x))
-    grad = prob.gradient(u)
-    assert np.allclose(grad, np.exp(u))
-
-    # Test jacobian
-    jac = prob.jacobian(u)
-    assert jac.shape == (7, 3)  # 3 + 3 + 1 rows
-
-
-def test_problem_multiple_variables():
-    """Test Problem with multiple CVXPY variables."""
-    x = cp.Variable(2)
-    y = cp.Variable(2)
-    obj = cp.sum(cp.log(x)) + cp.sum(cp.exp(y))
-    constraints = [
-        cp.log(x),
-        cp.exp(y),
-    ]
-
-    cvxpy_prob = cp.Problem(cp.Minimize(obj), constraints)
-    prob = Problem(cvxpy_prob)
-
-    x_vals = np.array([1.0, 2.0])
-    y_vals = np.array([0.5, 1.0])
-    u = np.concatenate([x_vals, y_vals])
-    prob.allocate(u)
-
-    # Test forward
-    obj_val, constraint_vals = prob.forward(u)
-    expected_obj = np.sum(np.log(x_vals)) + np.sum(np.exp(y_vals))
-    assert np.allclose(obj_val, expected_obj)
-    expected_constraint_vals = np.concatenate([np.log(x_vals), np.exp(y_vals)])
-    assert np.allclose(constraint_vals, expected_constraint_vals)
-
-    # Test gradient
-    grad = prob.gradient(u)
-    expected_grad = np.concatenate([1.0 / x_vals, np.exp(y_vals)])
-    assert np.allclose(grad, expected_grad)
-
-    # Test jacobian
-    jac = prob.jacobian(u)
-    assert jac.shape == (4, 4)
-    # First constraint log(x) only depends on x (first 2 vars)
-    # Second constraint exp(y) only depends on y (last 2 vars)
-    expected_jac = np.zeros((4, 4))
-    expected_jac[0, 0] = 1.0 / x_vals[0]
-    expected_jac[1, 1] = 1.0 / x_vals[1]
-    expected_jac[2, 2] = np.exp(y_vals[0])
-    expected_jac[3, 3] = np.exp(y_vals[1])
-    assert np.allclose(jac.toarray(), expected_jac)
-
-
-def test_problem_no_constraints_convert():
-    """Test Problem class with no constraints using convert."""
-    x = cp.Variable(4)
-    obj = cp.sum(cp.log(x))
-
-    cvxpy_prob = cp.Problem(cp.Minimize(obj))
-    prob = Problem(cvxpy_prob)
-
-    u = np.array([1.0, 2.0, 3.0, 4.0])
-    prob.allocate(u)
-
-    obj_val, constraint_vals = prob.forward(u)
-    assert np.allclose(obj_val, np.sum(np.log(u)))
-    assert len(constraint_vals) == 0
-
-    grad = prob.gradient(u)
-    assert np.allclose(grad, 1.0 / u)
-
-    jac = prob.jacobian(u)
-    assert jac.shape == (0, 4)
-
-
-def test_problem_larger_scale():
-    """Test Problem with larger variables and multiple constraints."""
-    n = 50
-    x = cp.Variable(n)
-    obj = cp.sum(cp.log(x))
-    constraints = [
-        cp.log(x),
-        cp.exp(x),
-        cp.sum(cp.log(x)),
-        cp.sum(cp.exp(x)),
-    ]
-
-    cvxpy_prob = cp.Problem(cp.Minimize(obj), constraints)
-    prob = Problem(cvxpy_prob)
-
-    u = np.linspace(1.0, 5.0, n)
-    prob.allocate(u)
-
-    # Test forward
-    obj_val, constraint_vals = prob.forward(u)
-    assert np.allclose(obj_val, np.sum(np.log(u)))
-    expected_constraint_vals = np.concatenate([
-        np.log(u),
-        np.exp(u),
-        [np.sum(np.log(u))],
-        [np.sum(np.exp(u))],
-    ])
-    assert np.allclose(constraint_vals, expected_constraint_vals)
-
-    # Test gradient
-    grad = prob.gradient(u)
-    assert np.allclose(grad, 1.0 / u)
-
-    # Test jacobian shape
-    jac = prob.jacobian(u)
-    assert jac.shape == (n + n + 1 + 1, n)  # 102 x 50
-
-
-def test_problem_repeated_evaluations():
-    """Test Problem with repeated evaluations at different points."""
-    x = cp.Variable(3)
-    obj = cp.sum(cp.log(x))
-    constraints = [cp.exp(x)]
-
-    cvxpy_prob = cp.Problem(cp.Minimize(obj), constraints)
-    prob = Problem(cvxpy_prob)
-
-    u1 = np.array([1.0, 2.0, 3.0])
-    prob.allocate(u1)
-
-    # First evaluation
-    obj_val1, _ = prob.forward(u1)
-    grad1 = prob.gradient(u1)
-
-    # Second evaluation at different point
-    u2 = np.array([2.0, 3.0, 4.0])
-    obj_val2, _ = prob.forward(u2)
-    grad2 = prob.gradient(u2)
-
-    assert np.allclose(obj_val1, np.sum(np.log(u1)))
-    assert np.allclose(obj_val2, np.sum(np.log(u2)))
-    assert np.allclose(grad1, 1.0 / u1)
-    assert np.allclose(grad2, 1.0 / u2)
-
-
-def test_problem_constraint_forward():
-    """Test Problem.constraint_forward for constraint values only."""
-    x = cp.Variable(2)
-    obj = cp.sum(cp.log(x))
-    constraints = [
-        cp.log(x),
-        cp.exp(x),
-    ]
-
-    cvxpy_prob = cp.Problem(cp.Minimize(obj), constraints)
-    prob = Problem(cvxpy_prob)
-
-    u = np.array([2.0, 4.0])
-    prob.allocate(u)
-
-    # Test constraint_forward
-    constraint_vals = prob.constraint_forward(u)
-    expected = np.concatenate([np.log(u), np.exp(u)])
-    assert np.allclose(constraint_vals, expected)
-
-    # Verify it gives same result as forward's constraint values
-    _, forward_constraint_vals = prob.forward(u)
-    assert np.allclose(constraint_vals, forward_constraint_vals)
-
-
-if __name__ == "__main__":
-    # Low-level tests
-    test_problem_forward_lowlevel()
-    test_problem_gradient_lowlevel()
-    test_problem_jacobian_lowlevel()
-    test_problem_constraint_forward_lowlevel()
-    test_problem_no_constraints_lowlevel()
-    # Problem class tests
-    test_problem_single_constraint()
-    test_problem_two_constraints()
-    test_problem_three_constraints_different_sizes()
-    test_problem_multiple_variables()
-    test_problem_no_constraints_convert()
-    test_problem_larger_scale()
-    test_problem_repeated_evaluations()
-    test_problem_constraint_forward()
-    print("All problem tests passed!")
diff --git a/src/affine/constant.c b/src/affine/constant.c
index d427ec9..577029b 100644
--- a/src/affine/constant.c
+++ b/src/affine/constant.c
@@ -8,6 +8,19 @@ static void forward(expr *node, const double *u)
     (void) u;
 }
 
+static void jacobian_init(expr *node)
+{
+    /* Constant jacobian is all zeros: size x n_vars with 0 nonzeros.
+     * new_csr_matrix uses calloc for row pointers, so they're already 0. */
+    node->jacobian = new_csr_matrix(node->size, node->n_vars, 0);
+}
+
+static void eval_jacobian(expr *node)
+{
+    /* Constant jacobian never changes - nothing to evaluate */
+    (void) node;
+}
+
 static bool is_affine(const expr *node)
 {
     (void) node;
@@ -20,6 +33,8 @@ expr *new_constant(int d1, int d2, int n_vars, const double *values)
     memcpy(node->value, values, d1 * d2 * sizeof(double));
     node->forward = forward;
     node->is_affine = is_affine;
+    node->jacobian_init = jacobian_init;
+    node->eval_jacobian = eval_jacobian;
 
     return node;
 }
diff --git a/src/affine/variable.c b/src/affine/variable.c
index 46e2446..d57bd96 100644
--- a/src/affine/variable.c
+++ b/src/affine/variable.c
@@ -20,6 +20,12 @@ static void jacobian_init(expr *node)
     node->jacobian->p[size] = size;
 }
 
+static void eval_jacobian(expr *node)
+{
+    /* Variable jacobian never changes - nothing to evaluate */
+    (void) node;
+}
+
 static bool is_affine(const expr *node)
 {
     (void) node;
@@ -33,6 +39,7 @@ expr *new_variable(int d1, int d2, int var_id, int n_vars)
     node->var_id = var_id;
     node->is_affine = is_affine;
     node->jacobian_init = jacobian_init;
+    node->eval_jacobian = eval_jacobian;
 
     return node;
 }
diff --git a/src/problem.c b/src/problem.c
index 45f0707..5fc9ad1 100644
--- a/src/problem.c
+++ b/src/problem.c
@@ -153,23 +153,11 @@ void free_problem(problem *prob)
     free(prob);
 }
 
-double problem_forward(problem *prob, const double *u)
+double problem_objective_forward(problem *prob, const double *u)
 {
-    /* Evaluate objective */
+    /* Evaluate objective only */
     prob->objective->forward(prob->objective, u);
-    double obj_val = prob->objective->value[0];
-
-    /* Evaluate constraints and copy values */
-    int offset = 0;
-    for (int i = 0; i < prob->n_constraints; i++)
-    {
-        expr *c = prob->constraints[i];
-        c->forward(c, u);
-        memcpy(prob->constraint_values + offset, c->value, c->size * sizeof(double));
-        offset += c->size;
-    }
-
-    return obj_val;
+    return prob->objective->value[0];
 }
 
 double *problem_constraint_forward(problem *prob, const double *u)
diff --git a/tests/all_tests.c b/tests/all_tests.c
index 751ad7c..66d7f62 100644
--- a/tests/all_tests.c
+++ b/tests/all_tests.c
@@ -6,6 +6,8 @@
 #include "forward_pass/affine/test_add.h"
 #include "forward_pass/affine/test_hstack.h"
 #include "forward_pass/affine/test_linear_op.h"
+#include "forward_pass/affine/test_neg.h"
+#include "forward_pass/affine/test_promote.h"
 #include "forward_pass/affine/test_sum.h"
 #include "forward_pass/affine/test_variable_constant.h"
 #include "forward_pass/composite/test_composite.h"
@@ -45,6 +47,12 @@ int main(void)
     mu_run_test(test_constant, tests_run);
     mu_run_test(test_addition, tests_run);
     mu_run_test(test_linear_op, tests_run);
+    mu_run_test(test_neg_forward, tests_run);
+    mu_run_test(test_neg_jacobian, tests_run);
+    mu_run_test(test_neg_chain, tests_run);
+    mu_run_test(test_promote_scalar_to_vector, tests_run);
+    mu_run_test(test_promote_scalar_jacobian, tests_run);
+    mu_run_test(test_promote_vector_jacobian, tests_run);
     mu_run_test(test_exp, tests_run);
     mu_run_test(test_log, tests_run);
     mu_run_test(test_composite, tests_run);
@@ -121,7 +129,7 @@ int main(void)
 
     printf("\n--- Problem Struct Tests ---\n");
     mu_run_test(test_problem_new_free, tests_run);
-    mu_run_test(test_problem_forward, tests_run);
+    mu_run_test(test_problem_objective_forward, tests_run);
     mu_run_test(test_problem_gradient, tests_run);
     mu_run_test(test_problem_jacobian, tests_run);
     mu_run_test(test_problem_jacobian_multi, tests_run);
diff --git a/tests/problem/test_problem.h b/tests/problem/test_problem.h
index 4ecc3d1..814a3c6 100644
--- a/tests/problem/test_problem.h
+++ b/tests/problem/test_problem.h
@@ -48,10 +48,10 @@ const char *test_problem_new_free(void)
 }
 
 /*
- * Test problem_forward: minimize sum(log(x))
+ * Test problem_objective_forward: minimize sum(log(x))
  *                       subject to x (as constraint)
  */
-const char *test_problem_forward(void)
+const char *test_problem_objective_forward(void)
 {
     expr *x = new_variable(3, 1, 0, 3);
     expr *log_x = new_log(x);
@@ -65,12 +65,15 @@ const char *test_problem_forward(void)
     double u[3] = {1.0, 2.0, 3.0};
     problem_allocate(prob, u);
 
-    double obj_val = problem_forward(prob, u);
+    double obj_val = problem_objective_forward(prob, u);
 
     /* Expected: sum(log([1, 2, 3])) = 0 + log(2) + log(3) */
     double expected_obj = log(1.0) + log(2.0) + log(3.0);
     mu_assert("objective value wrong", fabs(obj_val - expected_obj) < 1e-10);
 
+    /* Now evaluate constraints separately */
+    problem_constraint_forward(prob, u);
+
     /* Constraint values should be [1, 2, 3] */
     mu_assert("constraint[0] wrong", fabs(prob->constraint_values[0] - 1.0) < 1e-10);
     mu_assert("constraint[1] wrong", fabs(prob->constraint_values[1] - 2.0) < 1e-10);
@@ -85,6 +88,55 @@ const char *test_problem_forward(void)
     return 0;
 }
 
+/*
+ * Test problem_constraint_forward: evaluate constraints only
+ *   Constraint 1: log(x) -> [log(2), log(4)]
+ *   Constraint 2: exp(x) -> [exp(2), exp(4)]
+ */
+const char *test_problem_constraint_forward(void)
+{
+    int n_vars = 2;
+
+    /* Shared variable */
+    expr *x = new_variable(2, 1, 0, n_vars);
+
+    /* Objective: sum(log(x)) */
+    expr *log_obj = new_log(x);
+    expr *objective = new_sum(log_obj, -1);
+
+    /* Constraint 1: log(x) */
+    expr *log_c1 = new_log(x);
+
+    /* Constraint 2: exp(x) */
+    expr *exp_c2 = new_exp(x);
+
+    expr *constraints[2] = {log_c1, exp_c2};
+
+    problem *prob = new_problem(objective, constraints, 2);
+
+    double u[2] = {2.0, 4.0};
+    problem_allocate(prob, u);
+
+    double *constraint_vals = problem_constraint_forward(prob, u);
+
+    /* Check constraint values:
+     * [log(2), log(4), exp(2), exp(4)]
+     */
+    mu_assert("constraint[0] wrong", fabs(constraint_vals[0] - log(2.0)) < 1e-10);
+    mu_assert("constraint[1] wrong", fabs(constraint_vals[1] - log(4.0)) < 1e-10);
+    mu_assert("constraint[2] wrong", fabs(constraint_vals[2] - exp(2.0)) < 1e-10);
+    mu_assert("constraint[3] wrong", fabs(constraint_vals[3] - exp(4.0)) < 1e-10);
+
+    free_problem(prob);
+    free_expr(objective);
+    free_expr(log_obj);
+    free_expr(log_c1);
+    free_expr(exp_c2);
+    free_expr(x);
+
+    return 0;
+}
+
 /*
  * Test problem_gradient: gradient of sum(log(x)) = [1/x1, 1/x2, 1/x3]
  */
@@ -203,7 +255,6 @@ const char *test_problem_jacobian_multi(void)
 
     double u[2] = {2.0, 4.0};
     problem_allocate(prob, u);
-    problem_forward(prob, u);
 
     CSR_Matrix *jac = problem_jacobian(prob, u);
 
@@ -246,53 +297,4 @@ const char *test_problem_jacobian_multi(void)
     return 0;
 }
 
-/*
- * Test problem_constraint_forward: evaluate constraints only
- *   Constraint 1: log(x) -> [log(2), log(4)]
- *   Constraint 2: exp(x) -> [exp(2), exp(4)]
- */
-const char *test_problem_constraint_forward(void)
-{
-    int n_vars = 2;
-
-    /* Shared variable */
-    expr *x = new_variable(2, 1, 0, n_vars);
-
-    /* Objective: sum(log(x)) */
-    expr *log_obj = new_log(x);
-    expr *objective = new_sum(log_obj, -1);
-
-    /* Constraint 1: log(x) */
-    expr *log_c1 = new_log(x);
-
-    /* Constraint 2: exp(x) */
-    expr *exp_c2 = new_exp(x);
-
-    expr *constraints[2] = {log_c1, exp_c2};
-
-    problem *prob = new_problem(objective, constraints, 2);
-
-    double u[2] = {2.0, 4.0};
-    problem_allocate(prob, u);
-
-    double *constraint_vals = problem_constraint_forward(prob, u);
-
-    /* Check constraint values:
-     * [log(2), log(4), exp(2), exp(4)]
-     */
-    mu_assert("constraint[0] wrong", fabs(constraint_vals[0] - log(2.0)) < 1e-10);
-    mu_assert("constraint[1] wrong", fabs(constraint_vals[1] - log(4.0)) < 1e-10);
-    mu_assert("constraint[2] wrong", fabs(constraint_vals[2] - exp(2.0)) < 1e-10);
-    mu_assert("constraint[3] wrong", fabs(constraint_vals[3] - exp(4.0)) < 1e-10);
-
-    free_problem(prob);
-    free_expr(objective);
-    free_expr(log_obj);
-    free_expr(log_c1);
-    free_expr(exp_c2);
-    free_expr(x);
-
-    return 0;
-}
-
 #endif /* TEST_PROBLEM_H */

From 9a8cb0142eae6df73a69147fc5b9feada502005a Mon Sep 17 00:00:00 2001
From: William Zijie Zhang <william@gridmatic.com>
Date: Mon, 5 Jan 2026 21:05:13 -0500
Subject: [PATCH 08/27] adds new files

---
 python/.gitignore                        |   1 +
 python/tests/test_constrained.py         | 202 +++++++++++
 python/tests/test_problem_native.py      | 162 +++++++++
 python/tests/test_unconstrained.py       | 433 +++++++++++++++++++++++
 src/affine/neg.c                         | 114 ++++++
 src/affine/promote.c                     | 134 +++++++
 tests/forward_pass/affine/test_neg.h     |  80 +++++
 tests/forward_pass/affine/test_promote.h |  90 +++++
 8 files changed, 1216 insertions(+)
 create mode 100644 python/.gitignore
 create mode 100644 python/tests/test_constrained.py
 create mode 100644 python/tests/test_problem_native.py
 create mode 100644 python/tests/test_unconstrained.py
 create mode 100644 src/affine/neg.c
 create mode 100644 src/affine/promote.c
 create mode 100644 tests/forward_pass/affine/test_neg.h
 create mode 100644 tests/forward_pass/affine/test_promote.h

diff --git a/python/.gitignore b/python/.gitignore
new file mode 100644
index 0000000..ba0430d
--- /dev/null
+++ b/python/.gitignore
@@ -0,0 +1 @@
+__pycache__/
\ No newline at end of file
diff --git a/python/tests/test_constrained.py b/python/tests/test_constrained.py
new file mode 100644
index 0000000..23962f0
--- /dev/null
+++ b/python/tests/test_constrained.py
@@ -0,0 +1,202 @@
+import cvxpy as cp
+import numpy as np
+from convert import C_problem
+
+# Note: CVXPY converts constraints A >= B to B - A <= 0
+# So constr.expr for "log(x) >= 0" is "0 - log(x)" = -log(x)
+# All constraint values and jacobians are negated compared to the LHS
+
+
+def test_single_constraint():
+    """Test C_problem with single constraint."""
+    x = cp.Variable(3)
+    obj = cp.sum(cp.log(x))
+    constraints = [cp.log(x) >= 0]  # becomes 0 - log(x) <= 0
+
+    problem = cp.Problem(cp.Minimize(obj), constraints)
+    prob = C_problem(problem)
+
+    u = np.array([1.0, 2.0, 3.0])
+    prob.allocate(u)
+
+    # Test constraint_forward: constr.expr = -log(x)
+    constraint_vals = prob.constraint_forward(u)
+    assert np.allclose(constraint_vals, -np.log(u))
+
+    # Test jacobian: d/dx(-log(x)) = -1/x
+    jac = prob.jacobian(u)
+    expected_jac = np.diag(-1.0 / u)
+    assert np.allclose(jac.toarray(), expected_jac)
+
+
+def test_two_constraints():
+    """Test C_problem with two constraints."""
+    x = cp.Variable(2)
+    obj = cp.sum(cp.log(x))
+    constraints = [
+        cp.log(x) >= 0,  # becomes -log(x) <= 0
+        cp.exp(x) >= 0,  # becomes -exp(x) <= 0
+    ]
+
+    problem = cp.Problem(cp.Minimize(obj), constraints)
+    prob = C_problem(problem)
+
+    u = np.array([1.0, 2.0])
+    prob.allocate(u)
+
+    # Test constraint_forward: [-log(u), -exp(u)]
+    expected_constraint_vals = np.concatenate([-np.log(u), -np.exp(u)])
+    constraint_vals = prob.constraint_forward(u)
+    assert np.allclose(constraint_vals, expected_constraint_vals)
+
+    # Test jacobian - stacked vertically
+    jac = prob.jacobian(u)
+    assert jac.shape == (4, 2)
+    expected_jac = np.vstack([np.diag(-1.0 / u), np.diag(-np.exp(u))])
+    assert np.allclose(jac.toarray(), expected_jac)
+
+
+def test_three_constraints_different_sizes():
+    """Test C_problem with three constraints of different types."""
+    x = cp.Variable(3)
+    obj = cp.sum(cp.exp(x))
+    constraints = [
+        cp.log(x) >= 0,  # 3 constraints, becomes -log(x) <= 0
+        cp.exp(x) >= 0,  # 3 constraints, becomes -exp(x) <= 0
+        cp.sum(cp.log(x)) >= 0,  # 1 constraint, becomes -sum(log(x)) <= 0
+    ]
+
+    problem = cp.Problem(cp.Minimize(obj), constraints)
+    prob = C_problem(problem)
+
+    u = np.array([1.0, 2.0, 3.0])
+    prob.allocate(u)
+
+    # Test constraint_forward
+    expected_constraint_vals = np.concatenate([
+        -np.log(u),
+        -np.exp(u),
+        [-np.sum(np.log(u))]
+    ])
+    constraint_vals = prob.constraint_forward(u)
+    assert np.allclose(constraint_vals, expected_constraint_vals)
+
+    # Test jacobian shape and values
+    jac = prob.jacobian(u)
+    assert jac.shape == (7, 3)
+    # First 3 rows: -diag(1/u), next 3 rows: -diag(exp(u)), last row: -1/u
+    expected_jac = np.zeros((7, 3))
+    expected_jac[:3, :] = np.diag(-1.0 / u)
+    expected_jac[3:6, :] = np.diag(-np.exp(u))
+    expected_jac[6, :] = -1.0 / u
+    assert np.allclose(jac.toarray(), expected_jac)
+
+
+def test_multiple_variables():
+    """Test C_problem with multiple CVXPY variables."""
+    x = cp.Variable(2)
+    y = cp.Variable(2)
+    obj = cp.sum(cp.log(x)) + cp.sum(cp.exp(y))
+    constraints = [
+        cp.log(x) >= 0,  # becomes -log(x) <= 0
+        cp.exp(y) >= 0,  # becomes -exp(y) <= 0
+    ]
+
+    problem = cp.Problem(cp.Minimize(obj), constraints)
+    prob = C_problem(problem)
+
+    x_vals = np.array([1.0, 2.0])
+    y_vals = np.array([0.5, 1.0])
+    u = np.concatenate([x_vals, y_vals])
+    prob.allocate(u)
+
+    # Test constraint_forward
+    expected_constraint_vals = np.concatenate([-np.log(x_vals), -np.exp(y_vals)])
+    constraint_vals = prob.constraint_forward(u)
+    assert np.allclose(constraint_vals, expected_constraint_vals)
+
+    # Test jacobian
+    jac = prob.jacobian(u)
+    assert jac.shape == (4, 4)
+    expected_jac = np.zeros((4, 4))
+    expected_jac[0, 0] = -1.0 / x_vals[0]
+    expected_jac[1, 1] = -1.0 / x_vals[1]
+    expected_jac[2, 2] = -np.exp(y_vals[0])
+    expected_jac[3, 3] = -np.exp(y_vals[1])
+    assert np.allclose(jac.toarray(), expected_jac)
+
+
+def test_larger_scale():
+    """Test C_problem with larger variables and multiple constraints."""
+    n = 50
+    x = cp.Variable(n)
+    obj = cp.sum(cp.log(x))
+    constraints = [
+        cp.log(x) >= 0,  # n constraints
+        cp.exp(x) >= 0,  # n constraints
+        cp.sum(cp.log(x)) >= 0,  # 1 constraint
+        cp.sum(cp.exp(x)) >= 0,  # 1 constraint
+    ]
+
+    problem = cp.Problem(cp.Minimize(obj), constraints)
+    prob = C_problem(problem)
+
+    u = np.linspace(1.0, 5.0, n)
+    prob.allocate(u)
+
+    # Test constraint_forward
+    expected_constraint_vals = np.concatenate([
+        -np.log(u),
+        -np.exp(u),
+        [-np.sum(np.log(u))],
+        [-np.sum(np.exp(u))],
+    ])
+    constraint_vals = prob.constraint_forward(u)
+    assert np.allclose(constraint_vals, expected_constraint_vals)
+
+    # Test jacobian shape
+    jac = prob.jacobian(u)
+    assert jac.shape == (n + n + 1 + 1, n)
+
+
+def test_repeated_evaluations():
+    """Test C_problem with repeated evaluations at different points."""
+    x = cp.Variable(3)
+    obj = cp.sum(cp.log(x))
+    constraints = [cp.exp(x) >= 0]  # becomes -exp(x) <= 0
+
+    problem = cp.Problem(cp.Minimize(obj), constraints)
+    prob = C_problem(problem)
+
+    u1 = np.array([1.0, 2.0, 3.0])
+    prob.allocate(u1)
+
+    # First evaluation
+    constraint_vals1 = prob.constraint_forward(u1)
+    jac1 = prob.jacobian(u1)
+
+    # Second evaluation at different point
+    u2 = np.array([2.0, 3.0, 4.0])
+    constraint_vals2 = prob.constraint_forward(u2)
+    jac2 = prob.jacobian(u2)
+
+    assert np.allclose(constraint_vals1, -np.exp(u1))
+    assert np.allclose(constraint_vals2, -np.exp(u2))
+    assert np.allclose(jac1.toarray(), np.diag(-np.exp(u1)))
+    assert np.allclose(jac2.toarray(), np.diag(-np.exp(u2)))
+
+
+if __name__ == "__main__":
+    test_single_constraint()
+    print("test_single_constraint passed!")
+    test_two_constraints()
+    print("test_two_constraints passed!")
+    test_three_constraints_different_sizes()
+    print("test_three_constraints_different_sizes passed!")
+    test_multiple_variables()
+    print("test_multiple_variables passed!")
+    test_larger_scale()
+    print("test_larger_scale passed!")
+    test_repeated_evaluations()
+    print("test_repeated_evaluations passed!")
+    print("\nAll constrained tests passed!")
diff --git a/python/tests/test_problem_native.py b/python/tests/test_problem_native.py
new file mode 100644
index 0000000..a7ce1c4
--- /dev/null
+++ b/python/tests/test_problem_native.py
@@ -0,0 +1,162 @@
+import numpy as np
+from scipy import sparse
+import DNLP_diff_engine as diffengine
+
+
+def test_problem_objective_forward():
+    """Test problem_objective_forward and problem_constraint_forward (low-level)."""
+    n_vars = 3
+    x = diffengine.make_variable(n_vars, 1, 0, n_vars)
+    log_x = diffengine.make_log(x)
+    objective = diffengine.make_sum(log_x, -1)
+    constraints = [log_x]
+
+    prob = diffengine.make_problem(objective, constraints)
+    u = np.array([1.0, 2.0, 3.0])
+    diffengine.problem_allocate(prob, u)
+
+    obj_val = diffengine.problem_objective_forward(prob, u)
+    constraint_vals = diffengine.problem_constraint_forward(prob, u)
+
+    expected_obj = np.sum(np.log(u))
+    assert np.allclose(obj_val, expected_obj)
+    assert np.allclose(constraint_vals, np.log(u))
+
+
+def test_problem_constraint_forward():
+    """Test problem_constraint_forward for constraint values only (low-level)."""
+    n_vars = 2
+    x = diffengine.make_variable(n_vars, 1, 0, n_vars)
+
+    log_obj = diffengine.make_log(x)
+    objective = diffengine.make_sum(log_obj, -1)
+
+    log_c = diffengine.make_log(x)
+    exp_c = diffengine.make_exp(x)
+    constraints = [log_c, exp_c]
+
+    prob = diffengine.make_problem(objective, constraints)
+    u = np.array([2.0, 4.0])
+    diffengine.problem_allocate(prob, u)
+
+    constraint_vals = diffengine.problem_constraint_forward(prob, u)
+
+    # Expected: [log(2), log(4), exp(2), exp(4)]
+    expected = np.concatenate([np.log(u), np.exp(u)])
+    assert np.allclose(constraint_vals, expected)
+
+
+def test_problem_gradient():
+    """Test problem_gradient for objective gradient (low-level)."""
+    n_vars = 3
+    x = diffengine.make_variable(n_vars, 1, 0, n_vars)
+    log_x = diffengine.make_log(x)
+    objective = diffengine.make_sum(log_x, -1)
+
+    prob = diffengine.make_problem(objective, [])
+    u = np.array([1.0, 2.0, 4.0])
+    diffengine.problem_allocate(prob, u)
+
+    grad = diffengine.problem_gradient(prob, u)
+    expected_grad = 1.0 / u
+    assert np.allclose(grad, expected_grad)
+
+
+def test_problem_jacobian():
+    """Test problem_jacobian for constraint jacobian (low-level)."""
+    n_vars = 2
+    x = diffengine.make_variable(n_vars, 1, 0, n_vars)
+    log_x = diffengine.make_log(x)
+    objective = diffengine.make_sum(log_x, -1)
+    constraints = [log_x]
+
+    prob = diffengine.make_problem(objective, constraints)
+    u = np.array([2.0, 4.0])
+    diffengine.problem_allocate(prob, u)
+
+    data, indices, indptr, shape = diffengine.problem_jacobian(prob, u)
+    jac = sparse.csr_matrix((data, indices, indptr), shape=shape)
+
+    expected_jac = np.diag(1.0 / u)
+    assert np.allclose(jac.toarray(), expected_jac)
+
+
+def test_problem_no_constraints():
+    """Test Problem with no constraints (low-level)."""
+    n_vars = 3
+    x = diffengine.make_variable(n_vars, 1, 0, n_vars)
+    log_x = diffengine.make_log(x)
+    objective = diffengine.make_sum(log_x, -1)
+
+    prob = diffengine.make_problem(objective, [])
+    u = np.array([1.0, 2.0, 3.0])
+    diffengine.problem_allocate(prob, u)
+
+    obj_val = diffengine.problem_objective_forward(prob, u)
+    constraint_vals = diffengine.problem_constraint_forward(prob, u)
+    assert np.allclose(obj_val, np.sum(np.log(u)))
+    assert len(constraint_vals) == 0
+
+    grad = diffengine.problem_gradient(prob, u)
+    assert np.allclose(grad, 1.0 / u)
+
+    data, indices, indptr, shape = diffengine.problem_jacobian(prob, u)
+    jac = sparse.csr_matrix((data, indices, indptr), shape=shape)
+    assert jac.shape == (0, 3)
+
+
+def test_problem_multiple_constraints():
+    """Test problem with multiple different constraints (low-level)."""
+    n_vars = 3
+    x = diffengine.make_variable(n_vars, 1, 0, n_vars)
+    
+    # Objective: sum(log(x))
+    log_x = diffengine.make_log(x)
+    objective = diffengine.make_sum(log_x, -1)
+    
+    # Multiple constraints using the same variable:
+    # Constraint 1: log(x) - reused from objective
+    # Constraint 2: exp(x)
+    exp_x = diffengine.make_exp(x)
+    constraints = [log_x, exp_x]
+    
+    prob = diffengine.make_problem(objective, constraints)
+    u = np.array([1.0, 2.0, 3.0])
+    diffengine.problem_allocate(prob, u)
+    
+    # Test forward pass
+    obj_val = diffengine.problem_objective_forward(prob, u)
+    constraint_vals = diffengine.problem_constraint_forward(prob, u)
+    expected_obj = np.sum(np.log(u))
+    expected_constraints = np.concatenate([np.log(u), np.exp(u)])
+    assert np.allclose(obj_val, expected_obj)
+    assert np.allclose(constraint_vals, expected_constraints)
+    
+    # Test gradient
+    grad = diffengine.problem_gradient(prob, u)
+    expected_grad = 1.0 / u
+    assert np.allclose(grad, expected_grad)
+    
+    # Test Jacobian
+    data, indices, indptr, shape = diffengine.problem_jacobian(prob, u)
+    jac = sparse.csr_matrix((data, indices, indptr), shape=shape)
+    
+    # Expected Jacobian: 
+    # log(x): diag(1/u)
+    # exp(x): diag(exp(u))
+    expected_jac = np.vstack([
+        np.diag(1.0 / u),
+        np.diag(np.exp(u))
+    ])
+    assert jac.shape == (6, 3)
+    assert np.allclose(jac.toarray(), expected_jac)
+
+
+if __name__ == "__main__":
+    test_problem_objective_forward()
+    test_problem_constraint_forward()
+    test_problem_gradient()
+    test_problem_jacobian()
+    test_problem_no_constraints()
+    test_problem_multiple_constraints()
+    print("All native problem tests passed!")
diff --git a/python/tests/test_unconstrained.py b/python/tests/test_unconstrained.py
new file mode 100644
index 0000000..6cacd0b
--- /dev/null
+++ b/python/tests/test_unconstrained.py
@@ -0,0 +1,433 @@
+import cvxpy as cp
+import numpy as np
+from convert import C_problem
+
+# Note: Current implementation supports:
+# - Elementwise ops on variables: log(x), exp(x)
+# - Add of elementwise results: log(x) + exp(y)
+# - Sum reductions: sum(log(x))
+# - Multiple variables with separate operations
+#
+# NOT YET SUPPORTED (see unsupported section at bottom of file):
+# - Elementwise ops on add results: log(x + y)
+
+
+def test_sum_log():
+    """Test sum(log(x)) objective and gradient."""
+    x = cp.Variable(4)
+    problem = cp.Problem(cp.Minimize(cp.sum(cp.log(x))))
+    prob = C_problem(problem)
+
+    u = np.array([1.0, 2.0, 3.0, 4.0])
+    prob.allocate(u)
+
+    # Objective
+    obj_val = prob.objective_forward(u)
+    expected = np.sum(np.log(u))
+    assert np.allclose(obj_val, expected)
+
+    # Gradient: d/dx sum(log(x)) = 1/x
+    grad = prob.gradient(u)
+    expected_grad = 1.0 / u
+    assert np.allclose(grad, expected_grad)
+
+
+def test_sum_exp():
+    """Test sum(exp(x)) objective and gradient."""
+    x = cp.Variable(3)
+    problem = cp.Problem(cp.Minimize(cp.sum(cp.exp(x))))
+    prob = C_problem(problem)
+
+    u = np.array([0.0, 1.0, 2.0])
+    prob.allocate(u)
+
+    # Objective
+    obj_val = prob.objective_forward(u)
+    expected = np.sum(np.exp(u))
+    assert np.allclose(obj_val, expected)
+
+    # Gradient: d/dx sum(exp(x)) = exp(x)
+    grad = prob.gradient(u)
+    expected_grad = np.exp(u)
+    assert np.allclose(grad, expected_grad)
+
+
+def test_variable_reuse():
+    """Test sum(log(x) + exp(x)) - same variable used twice."""
+    x = cp.Variable(2)
+    problem = cp.Problem(cp.Minimize(cp.sum(cp.log(x) + cp.exp(x))))
+    prob = C_problem(problem)
+
+    u = np.array([1.0, 2.0])
+    prob.allocate(u)
+
+    # Objective
+    obj_val = prob.objective_forward(u)
+    expected = np.sum(np.log(u) + np.exp(u))
+    assert np.allclose(obj_val, expected)
+
+    # Gradient: d/dx_i = 1/x_i + exp(x_i)
+    grad = prob.gradient(u)
+    expected_grad = 1.0 / u + np.exp(u)
+    assert np.allclose(grad, expected_grad)
+
+
+def test_variable_used_multiple_times():
+    """Test sum(log(x) + exp(x) + log(x)) - variable used 3 times."""
+    x = cp.Variable(3)
+    problem = cp.Problem(cp.Minimize(cp.sum(cp.log(x) + cp.exp(x) + cp.log(x))))
+    prob = C_problem(problem)
+
+    u = np.array([1.0, 2.0, 3.0])
+    prob.allocate(u)
+
+    # Objective
+    obj_val = prob.objective_forward(u)
+    expected = np.sum(2 * np.log(u) + np.exp(u))
+    assert np.allclose(obj_val, expected)
+
+    # Gradient: d/dx_i = 2/x_i + exp(x_i)
+    grad = prob.gradient(u)
+    expected_grad = 2.0 / u + np.exp(u)
+    assert np.allclose(grad, expected_grad)
+
+
+def test_larger_variable():
+    """Test sum(log(x)) with larger variable (100 elements)."""
+    x = cp.Variable(100)
+    problem = cp.Problem(cp.Minimize(cp.sum(cp.log(x))))
+    prob = C_problem(problem)
+
+    u = np.linspace(1.0, 10.0, 100)
+    prob.allocate(u)
+
+    # Objective
+    obj_val = prob.objective_forward(u)
+    expected = np.sum(np.log(u))
+    assert np.allclose(obj_val, expected)
+
+    # Gradient
+    grad = prob.gradient(u)
+    expected_grad = 1.0 / u
+    assert np.allclose(grad, expected_grad)
+
+
+def test_matrix_variable():
+    """Test sum(log(X)) with 2D matrix variable (3x4)."""
+    X = cp.Variable((3, 4))
+    problem = cp.Problem(cp.Minimize(cp.sum(cp.log(X))))
+    prob = C_problem(problem)
+
+    u = np.arange(1.0, 13.0)  # 12 elements
+    prob.allocate(u)
+
+    # Objective
+    obj_val = prob.objective_forward(u)
+    expected = np.sum(np.log(u))
+    assert np.allclose(obj_val, expected)
+
+    # Gradient
+    grad = prob.gradient(u)
+    expected_grad = 1.0 / u
+    assert np.allclose(grad, expected_grad)
+
+
+def test_two_variables_separate_ops():
+    """Test sum(log(x)) + sum(exp(y)) - two variables with separate ops."""
+    x = cp.Variable(2)
+    y = cp.Variable(2)
+    problem = cp.Problem(cp.Minimize(cp.sum(cp.log(x)) + cp.sum(cp.exp(y))))
+    prob = C_problem(problem)
+
+    x_vals = np.array([1.0, 2.0])
+    y_vals = np.array([0.5, 1.0])
+    u = np.concatenate([x_vals, y_vals])
+    prob.allocate(u)
+
+    # Objective
+    obj_val = prob.objective_forward(u)
+    expected = np.sum(np.log(x_vals)) + np.sum(np.exp(y_vals))
+    assert np.allclose(obj_val, expected)
+
+    # Gradient
+    grad = prob.gradient(u)
+    expected_grad = np.concatenate([1.0 / x_vals, np.exp(y_vals)])
+    assert np.allclose(grad, expected_grad)
+
+
+def test_two_variables_same_sum():
+    """Test sum(log(x) + log(y)) - two variables added before sum."""
+    x = cp.Variable(2)
+    y = cp.Variable(2)
+    problem = cp.Problem(cp.Minimize(cp.sum(cp.log(x) + cp.log(y))))
+    prob = C_problem(problem)
+
+    x_vals = np.array([1.0, 2.0])
+    y_vals = np.array([3.0, 4.0])
+    u = np.concatenate([x_vals, y_vals])
+    prob.allocate(u)
+
+    # Objective
+    obj_val = prob.objective_forward(u)
+    expected = np.sum(np.log(x_vals) + np.log(y_vals))
+    assert np.allclose(obj_val, expected)
+
+    # Gradient
+    grad = prob.gradient(u)
+    expected_grad = np.concatenate([1.0 / x_vals, 1.0 / y_vals])
+    assert np.allclose(grad, expected_grad)
+
+
+def test_mixed_sizes():
+    """Test sum(log(a)) + sum(log(b)) + sum(log(c)) with different sized variables."""
+    a = cp.Variable(2)
+    b = cp.Variable(5)
+    c = cp.Variable(3)
+    problem = cp.Problem(cp.Minimize(
+        cp.sum(cp.log(a)) + cp.sum(cp.log(b)) + cp.sum(cp.log(c))
+    ))
+    prob = C_problem(problem)
+
+    a_vals = np.array([1.0, 2.0])
+    b_vals = np.array([1.0, 2.0, 3.0, 4.0, 5.0])
+    c_vals = np.array([1.0, 2.0, 3.0])
+    u = np.concatenate([a_vals, b_vals, c_vals])
+    prob.allocate(u)
+
+    # Objective
+    obj_val = prob.objective_forward(u)
+    expected = np.sum(np.log(a_vals)) + np.sum(np.log(b_vals)) + np.sum(np.log(c_vals))
+    assert np.allclose(obj_val, expected)
+
+    # Gradient
+    grad = prob.gradient(u)
+    expected_grad = 1.0 / u
+    assert np.allclose(grad, expected_grad)
+
+
+def test_multiple_variables_log_exp():
+    """Test sum(log(a)) + sum(log(b)) + sum(exp(c)) + sum(exp(d))."""
+    a = cp.Variable(2)
+    b = cp.Variable(2)
+    c = cp.Variable(2)
+    d = cp.Variable(2)
+    obj = cp.sum(cp.log(a)) + cp.sum(cp.log(b)) + cp.sum(cp.exp(c)) + cp.sum(cp.exp(d))
+    problem = cp.Problem(cp.Minimize(obj))
+    prob = C_problem(problem)
+
+    a_vals = np.array([1.0, 2.0])
+    b_vals = np.array([0.5, 1.0])
+    c_vals = np.array([0.1, 0.2])
+    d_vals = np.array([0.1, 0.1])
+    u = np.concatenate([a_vals, b_vals, c_vals, d_vals])
+    prob.allocate(u)
+
+    # Objective
+    obj_val = prob.objective_forward(u)
+    expected = (np.sum(np.log(a_vals)) + np.sum(np.log(b_vals)) +
+                np.sum(np.exp(c_vals)) + np.sum(np.exp(d_vals)))
+    assert np.allclose(obj_val, expected)
+
+    # Gradient
+    grad = prob.gradient(u)
+    expected_grad = np.concatenate([
+        1.0 / a_vals,
+        1.0 / b_vals,
+        np.exp(c_vals),
+        np.exp(d_vals)
+    ])
+    assert np.allclose(grad, expected_grad)
+
+
+def test_three_variables_mixed():
+    """Test sum(log(x) + exp(y) + log(z)) - three variables mixed."""
+    x = cp.Variable(2)
+    y = cp.Variable(2)
+    z = cp.Variable(2)
+    obj = cp.sum(cp.log(x) + cp.exp(y) + cp.log(z))
+    problem = cp.Problem(cp.Minimize(obj))
+    prob = C_problem(problem)
+
+    x_vals = np.array([1.0, 2.0])
+    y_vals = np.array([0.5, 1.0])
+    z_vals = np.array([2.0, 3.0])
+    u = np.concatenate([x_vals, y_vals, z_vals])
+    prob.allocate(u)
+
+    # Objective
+    obj_val = prob.objective_forward(u)
+    expected = np.sum(np.log(x_vals) + np.exp(y_vals) + np.log(z_vals))
+    assert np.allclose(obj_val, expected)
+
+    # Gradient
+    grad = prob.gradient(u)
+    expected_grad = np.concatenate([1.0 / x_vals, np.exp(y_vals), 1.0 / z_vals])
+    assert np.allclose(grad, expected_grad)
+
+
+def test_repeated_evaluations():
+    """Test repeated evaluations at different points."""
+    x = cp.Variable(3)
+    problem = cp.Problem(cp.Minimize(cp.sum(cp.log(x))))
+    prob = C_problem(problem)
+
+    u1 = np.array([1.0, 2.0, 3.0])
+    prob.allocate(u1)
+
+    # First evaluation
+    obj1 = prob.objective_forward(u1)
+    grad1 = prob.gradient(u1)
+
+    # Second evaluation
+    u2 = np.array([2.0, 3.0, 4.0])
+    obj2 = prob.objective_forward(u2)
+    grad2 = prob.gradient(u2)
+
+    assert np.allclose(obj1, np.sum(np.log(u1)))
+    assert np.allclose(obj2, np.sum(np.log(u2)))
+    assert np.allclose(grad1, 1.0 / u1)
+    assert np.allclose(grad2, 1.0 / u2)
+
+
+# =============================================================================
+# UNSUPPORTED PATTERNS
+# =============================================================================
+# The following patterns are NOT YET SUPPORTED because the current
+# implementation of elementwise univariate ops (log, exp, etc.) in
+# src/elementwise_univariate/common.c assumes that the child expression
+# is either a Variable or a LinearOp (which includes negation as -I).
+#
+# When the child is an Add expression (like x + y), the code incorrectly
+# casts it to linear_op_expr and reads invalid memory, causing crashes.
+#
+# To support these patterns, we would need to generalize the jacobian
+# computation in elementwise ops to handle arbitrary child expressions
+# by using the child's jacobian directly rather than assuming specific
+# structure.
+# =============================================================================
+
+# def test_elementwise_on_add():
+#     """NOT SUPPORTED: log(x + y) - elementwise op on add result.
+#
+#     This fails because log's jacobian computation assumes its child
+#     is a Variable or LinearOp, but here it's an Add expression.
+#     """
+#     x = cp.Variable(2)
+#     y = cp.Variable(2)
+#     problem = cp.Problem(cp.Minimize(cp.sum(cp.log(x + y))))
+#     prob = C_problem(problem)
+#
+#     x_vals = np.array([1.0, 2.0])
+#     y_vals = np.array([0.5, 1.0])
+#     u = np.concatenate([x_vals, y_vals])
+#     prob.allocate(u)
+#
+#     # Expected objective: sum(log(x + y))
+#     obj_val = prob.objective_forward(u)
+#     expected = np.sum(np.log(x_vals + y_vals))
+#     assert np.allclose(obj_val, expected)
+#
+#     # Expected gradient: d/dx_i = 1/(x_i + y_i), d/dy_i = 1/(x_i + y_i)
+#     grad = prob.gradient(u)
+#     grad_xy = 1.0 / (x_vals + y_vals)
+#     expected_grad = np.concatenate([grad_xy, grad_xy])
+#     assert np.allclose(grad, expected_grad)
+
+
+# def test_log_of_sum():
+#     """NOT SUPPORTED: log(sum(x)) - elementwise op on sum result.
+#
+#     This fails because log's jacobian computation assumes its child
+#     is a Variable or LinearOp, but here it's a Sum expression.
+#     """
+#     x = cp.Variable(4)
+#     problem = cp.Problem(cp.Minimize(cp.log(cp.sum(x))))
+#     prob = C_problem(problem)
+#
+#     u = np.array([1.0, 2.0, 3.0, 4.0])
+#     prob.allocate(u)
+#
+#     # Expected objective: log(sum(x))
+#     obj_val = prob.objective_forward(u)
+#     expected = np.log(np.sum(u))
+#     assert np.allclose(obj_val, expected)
+#
+#     # Expected gradient: d/dx_i log(sum(x)) = 1/sum(x)
+#     grad = prob.gradient(u)
+#     expected_grad = np.ones_like(u) / np.sum(u)
+#     assert np.allclose(grad, expected_grad)
+
+
+# def test_mixed_elementwise_on_add():
+#     """NOT SUPPORTED: log(x + y) + exp(y + z) - multiple elementwise on adds.
+#
+#     Same underlying issue - elementwise ops don't support Add children.
+#     """
+#     x = cp.Variable(2)
+#     y = cp.Variable(2)
+#     z = cp.Variable(2)
+#     obj = cp.sum(cp.log(x + y)) + cp.sum(cp.exp(y + z))
+#     problem = cp.Problem(cp.Minimize(obj))
+#     prob = C_problem(problem)
+#
+#     x_vals = np.array([1.0, 2.0])
+#     y_vals = np.array([0.5, 1.0])
+#     z_vals = np.array([0.1, 0.2])
+#     u = np.concatenate([x_vals, y_vals, z_vals])
+#     prob.allocate(u)
+#
+#     # Expected objective
+#     obj_val = prob.objective_forward(u)
+#     expected = np.sum(np.log(x_vals + y_vals)) + np.sum(np.exp(y_vals + z_vals))
+#     assert np.allclose(obj_val, expected)
+
+
+# def test_nested_add_in_elementwise():
+#     """NOT SUPPORTED: log(x + y + z) - nested adds inside elementwise.
+#
+#     Same underlying issue.
+#     """
+#     x = cp.Variable(2)
+#     y = cp.Variable(2)
+#     z = cp.Variable(2)
+#     problem = cp.Problem(cp.Minimize(cp.sum(cp.log(x + y + z))))
+#     prob = C_problem(problem)
+#
+#     x_vals = np.array([1.0, 2.0])
+#     y_vals = np.array([0.5, 1.0])
+#     z_vals = np.array([0.5, 0.5])
+#     u = np.concatenate([x_vals, y_vals, z_vals])
+#     prob.allocate(u)
+#
+#     # Expected objective
+#     obj_val = prob.objective_forward(u)
+#     expected = np.sum(np.log(x_vals + y_vals + z_vals))
+#     assert np.allclose(obj_val, expected)
+
+
+if __name__ == "__main__":
+    test_sum_log()
+    print("test_sum_log passed!")
+    test_sum_exp()
+    print("test_sum_exp passed!")
+    test_variable_reuse()
+    print("test_variable_reuse passed!")
+    test_variable_used_multiple_times()
+    print("test_variable_used_multiple_times passed!")
+    test_larger_variable()
+    print("test_larger_variable passed!")
+    test_matrix_variable()
+    print("test_matrix_variable passed!")
+    test_two_variables_separate_ops()
+    print("test_two_variables_separate_ops passed!")
+    test_two_variables_same_sum()
+    print("test_two_variables_same_sum passed!")
+    test_mixed_sizes()
+    print("test_mixed_sizes passed!")
+    test_multiple_variables_log_exp()
+    print("test_multiple_variables_log_exp passed!")
+    test_three_variables_mixed()
+    print("test_three_variables_mixed passed!")
+    test_repeated_evaluations()
+    print("test_repeated_evaluations passed!")
+    print("\nAll unconstrained tests passed!")
diff --git a/src/affine/neg.c b/src/affine/neg.c
new file mode 100644
index 0000000..0af70ce
--- /dev/null
+++ b/src/affine/neg.c
@@ -0,0 +1,114 @@
+#include "affine.h"
+#include <stdlib.h>
+
+static void forward(expr *node, const double *u)
+{
+    /* child's forward pass */
+    node->left->forward(node->left, u);
+
+    /* negate values */
+    for (int i = 0; i < node->size; i++)
+    {
+        node->value[i] = -node->left->value[i];
+    }
+}
+
+static void jacobian_init(expr *node)
+{
+    /* initialize child's jacobian */
+    node->left->jacobian_init(node->left);
+
+    /* 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);
+}
+
+static void eval_jacobian(expr *node)
+{
+    /* evaluate child's jacobian */
+    node->left->eval_jacobian(node->left);
+
+    /* negate child's jacobian: copy structure, negate values */
+    CSR_Matrix *child_jac = node->left->jacobian;
+    CSR_Matrix *jac = node->jacobian;
+
+    /* copy row pointers */
+    for (int i = 0; i <= child_jac->m; i++)
+    {
+        jac->p[i] = child_jac->p[i];
+    }
+
+    /* copy column indices and negate values */
+    int nnz = child_jac->p[child_jac->m];
+    for (int k = 0; k < nnz; k++)
+    {
+        jac->i[k] = child_jac->i[k];
+        jac->x[k] = -child_jac->x[k];
+    }
+    jac->nnz = nnz;
+}
+
+static void wsum_hess_init(expr *node)
+{
+    /* initialize child's wsum_hess */
+    node->left->wsum_hess_init(node->left);
+
+    /* 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);
+}
+
+static void eval_wsum_hess(expr *node, const double *w)
+{
+    /* For neg(x), d^2(-x)/dx^2 = 0, but we need to pass -w to child
+     * Actually: d/dx(-x) = -I, so Hessian contribution is (-I)^T H (-I) = H
+     * But the weight vector needs to be passed through: child sees same w */
+
+    /* Negate weights for child (chain rule for linear transformation) */
+    double *neg_w = (double *) malloc(node->size * sizeof(double));
+    for (int i = 0; i < node->size; i++)
+    {
+        neg_w[i] = -w[i];
+    }
+
+    /* evaluate child's wsum_hess with negated weights */
+    node->left->eval_wsum_hess(node->left, neg_w);
+    free(neg_w);
+
+    /* copy child's wsum_hess (the negation is already accounted for in weights) */
+    CSR_Matrix *child_hess = node->left->wsum_hess;
+    CSR_Matrix *hess = node->wsum_hess;
+
+    for (int i = 0; i <= child_hess->m; i++)
+    {
+        hess->p[i] = child_hess->p[i];
+    }
+
+    int nnz = child_hess->p[child_hess->m];
+    for (int k = 0; k < nnz; k++)
+    {
+        hess->i[k] = child_hess->i[k];
+        hess->x[k] = child_hess->x[k];
+    }
+    hess->nnz = nnz;
+}
+
+static bool is_affine(const expr *node)
+{
+    return node->left->is_affine(node->left);
+}
+
+expr *new_neg(expr *child)
+{
+    expr *node = new_expr(child->d1, child->d2, child->n_vars);
+    node->left = child;
+    expr_retain(child);
+    node->forward = forward;
+    node->is_affine = is_affine;
+    node->jacobian_init = jacobian_init;
+    node->eval_jacobian = eval_jacobian;
+    node->wsum_hess_init = wsum_hess_init;
+    node->eval_wsum_hess = eval_wsum_hess;
+
+    return node;
+}
diff --git a/src/affine/promote.c b/src/affine/promote.c
new file mode 100644
index 0000000..fcefc13
--- /dev/null
+++ b/src/affine/promote.c
@@ -0,0 +1,134 @@
+#include "affine.h"
+#include "subexpr.h"
+#include <stdlib.h>
+
+/* Promote broadcasts a child expression to a larger shape.
+ * Typically used to broadcast a scalar to a vector. */
+
+static void forward(expr *node, const double *u)
+{
+    promote_expr *prom = (promote_expr *) node;
+
+    /* child's forward pass */
+    node->left->forward(node->left, u);
+
+    /* broadcast child value to output shape */
+    int child_size = node->left->size;
+    for (int i = 0; i < node->size; i++)
+    {
+        /* replicate pattern: output[i] = child[i % child_size] */
+        node->value[i] = node->left->value[i % child_size];
+    }
+    (void) prom; /* unused for now, shape info stored in d1/d2 */
+}
+
+static void jacobian_init(expr *node)
+{
+    /* initialize child's jacobian */
+    node->left->jacobian_init(node->left);
+
+    /* Each output row copies a row from child's jacobian (with wrapping).
+     * For scalar->vector: all rows are copies of the single child row.
+     * nnz = (output_size / child_size) * child_jac_nnz */
+    CSR_Matrix *child_jac = node->left->jacobian;
+    int child_size = node->left->size;
+    int repeat = (node->size + child_size - 1) / child_size;
+    int nnz_max = repeat * child_jac->nnz;
+    if (nnz_max == 0) nnz_max = 1;
+
+    node->jacobian = new_csr_matrix(node->size, node->n_vars, nnz_max);
+}
+
+static void eval_jacobian(expr *node)
+{
+    /* evaluate child's jacobian */
+    node->left->eval_jacobian(node->left);
+
+    CSR_Matrix *child_jac = node->left->jacobian;
+    CSR_Matrix *jac = node->jacobian;
+    int child_size = node->left->size;
+
+    /* Build jacobian by replicating child's rows */
+    jac->nnz = 0;
+    for (int row = 0; row < node->size; row++)
+    {
+        jac->p[row] = jac->nnz;
+
+        /* which child row does this output row correspond to? */
+        int child_row = row % child_size;
+
+        /* copy entries from child_jac row */
+        for (int k = child_jac->p[child_row]; k < child_jac->p[child_row + 1]; k++)
+        {
+            jac->i[jac->nnz] = child_jac->i[k];
+            jac->x[jac->nnz] = child_jac->x[k];
+            jac->nnz++;
+        }
+    }
+    jac->p[node->size] = jac->nnz;
+}
+
+static void wsum_hess_init(expr *node)
+{
+    /* initialize child's wsum_hess */
+    node->left->wsum_hess_init(node->left);
+
+    /* same sparsity as child since we're summing weights */
+    CSR_Matrix *child_hess = node->left->wsum_hess;
+    node->wsum_hess = new_csr_matrix(child_hess->m, child_hess->n, child_hess->nnz);
+}
+
+static void eval_wsum_hess(expr *node, const double *w)
+{
+    /* Sum weights that correspond to the same child element */
+    int child_size = node->left->size;
+    double *summed_w = (double *) calloc(child_size, sizeof(double));
+
+    for (int i = 0; i < node->size; i++)
+    {
+        summed_w[i % child_size] += w[i];
+    }
+
+    /* evaluate child's wsum_hess with summed weights */
+    node->left->eval_wsum_hess(node->left, summed_w);
+    free(summed_w);
+
+    /* copy child's wsum_hess */
+    CSR_Matrix *child_hess = node->left->wsum_hess;
+    CSR_Matrix *hess = node->wsum_hess;
+
+    for (int i = 0; i <= child_hess->m; i++)
+    {
+        hess->p[i] = child_hess->p[i];
+    }
+
+    int nnz = child_hess->p[child_hess->m];
+    for (int k = 0; k < nnz; k++)
+    {
+        hess->i[k] = child_hess->i[k];
+        hess->x[k] = child_hess->x[k];
+    }
+    hess->nnz = nnz;
+}
+
+static bool is_affine(const expr *node)
+{
+    return node->left->is_affine(node->left);
+}
+
+expr *new_promote(expr *child, int d1, int d2)
+{
+    /* Allocate the type-specific struct */
+    promote_expr *prom = (promote_expr *) calloc(1, sizeof(promote_expr));
+    expr *node = &prom->base;
+
+    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;
+
+    return node;
+}
diff --git a/tests/forward_pass/affine/test_neg.h b/tests/forward_pass/affine/test_neg.h
new file mode 100644
index 0000000..73c14e3
--- /dev/null
+++ b/tests/forward_pass/affine/test_neg.h
@@ -0,0 +1,80 @@
+#include <math.h>
+#include <stdio.h>
+#include <stdlib.h>
+
+#include "affine.h"
+#include "expr.h"
+#include "minunit.h"
+#include "test_helpers.h"
+
+const char *test_neg_forward(void)
+{
+    double u[3] = {1.0, 2.0, 3.0};
+    expr *var = new_variable(3, 1, 0, 3);
+    expr *neg_node = new_neg(var);
+    neg_node->forward(neg_node, u);
+    double expected[3] = {-1.0, -2.0, -3.0};
+    mu_assert("neg forward failed", cmp_double_array(neg_node->value, expected, 3));
+    free_expr(neg_node);
+    free_expr(var);
+    return 0;
+}
+
+const char *test_neg_jacobian(void)
+{
+    double u[3] = {1.0, 2.0, 3.0};
+    expr *var = new_variable(3, 1, 0, 3);
+    expr *neg_node = new_neg(var);
+    neg_node->forward(neg_node, u);
+    neg_node->jacobian_init(neg_node);
+    neg_node->eval_jacobian(neg_node);
+
+    /* Jacobian of neg(x) is -I (diagonal with -1) */
+    double expected_x[3] = {-1.0, -1.0, -1.0};
+    int expected_p[4] = {0, 1, 2, 3};
+    int expected_i[3] = {0, 1, 2};
+
+    mu_assert("neg jacobian vals fail",
+              cmp_double_array(neg_node->jacobian->x, expected_x, 3));
+    mu_assert("neg jacobian rows fail",
+              cmp_int_array(neg_node->jacobian->p, expected_p, 4));
+    mu_assert("neg jacobian cols fail",
+              cmp_int_array(neg_node->jacobian->i, expected_i, 3));
+
+    free_expr(neg_node);
+    free_expr(var);
+    return 0;
+}
+
+const char *test_neg_chain(void)
+{
+    /* Test neg(neg(x)) = x */
+    double u[3] = {1.0, 2.0, 3.0};
+    expr *var = new_variable(3, 1, 0, 3);
+    expr *neg1 = new_neg(var);
+    expr *neg2 = new_neg(neg1);
+    neg2->forward(neg2, u);
+
+    /* neg(neg(x)) should equal x */
+    mu_assert("neg chain forward failed", cmp_double_array(neg2->value, u, 3));
+
+    neg2->jacobian_init(neg2);
+    neg2->eval_jacobian(neg2);
+
+    /* Jacobian of neg(neg(x)) is (-1)*(-1)*I = I */
+    double expected_x[3] = {1.0, 1.0, 1.0};
+    int expected_p[4] = {0, 1, 2, 3};
+    int expected_i[3] = {0, 1, 2};
+
+    mu_assert("neg chain jacobian vals fail",
+              cmp_double_array(neg2->jacobian->x, expected_x, 3));
+    mu_assert("neg chain jacobian rows fail",
+              cmp_int_array(neg2->jacobian->p, expected_p, 4));
+    mu_assert("neg chain jacobian cols fail",
+              cmp_int_array(neg2->jacobian->i, expected_i, 3));
+
+    free_expr(neg2);
+    free_expr(neg1);
+    free_expr(var);
+    return 0;
+}
diff --git a/tests/forward_pass/affine/test_promote.h b/tests/forward_pass/affine/test_promote.h
new file mode 100644
index 0000000..eda5bf3
--- /dev/null
+++ b/tests/forward_pass/affine/test_promote.h
@@ -0,0 +1,90 @@
+#include <math.h>
+#include <stdio.h>
+#include <stdlib.h>
+
+#include "affine.h"
+#include "expr.h"
+#include "minunit.h"
+#include "test_helpers.h"
+
+const char *test_promote_scalar_to_vector(void)
+{
+    /* Promote scalar to 4-element vector */
+    double u[1] = {5.0};
+    expr *var = new_variable(1, 1, 0, 1);
+    expr *promote_node = new_promote(var, 4, 1);
+    promote_node->forward(promote_node, u);
+
+    double expected[4] = {5.0, 5.0, 5.0, 5.0};
+    mu_assert("promote scalar->vector forward failed",
+              cmp_double_array(promote_node->value, expected, 4));
+
+    free_expr(promote_node);
+    free_expr(var);
+    return 0;
+}
+
+const char *test_promote_scalar_jacobian(void)
+{
+    /* Promote scalar to 3-element vector, check jacobian */
+    double u[1] = {2.0};
+    expr *var = new_variable(1, 1, 0, 1);
+    expr *promote_node = new_promote(var, 3, 1);
+    promote_node->forward(promote_node, u);
+    promote_node->jacobian_init(promote_node);
+    promote_node->eval_jacobian(promote_node);
+
+    /* Jacobian is 3x1 with all 1s (each output depends on same input) */
+    double expected_x[3] = {1.0, 1.0, 1.0};
+    int expected_p[4] = {0, 1, 2, 3};
+    int expected_i[3] = {0, 0, 0};
+
+    mu_assert("promote jacobian vals fail",
+              cmp_double_array(promote_node->jacobian->x, expected_x, 3));
+    mu_assert("promote jacobian rows fail",
+              cmp_int_array(promote_node->jacobian->p, expected_p, 4));
+    mu_assert("promote jacobian cols fail",
+              cmp_int_array(promote_node->jacobian->i, expected_i, 3));
+
+    free_expr(promote_node);
+    free_expr(var);
+    return 0;
+}
+
+const char *test_promote_vector_jacobian(void)
+{
+    /* Promote 2-vector to 4-vector, check jacobian */
+    double u[2] = {1.0, 2.0};
+    expr *var = new_variable(2, 1, 0, 2);
+    expr *promote_node = new_promote(var, 4, 1);
+    promote_node->forward(promote_node, u);
+
+    /* Pattern repeats: [1, 2, 1, 2] */
+    double expected_val[4] = {1.0, 2.0, 1.0, 2.0};
+    mu_assert("promote vector forward failed",
+              cmp_double_array(promote_node->value, expected_val, 4));
+
+    promote_node->jacobian_init(promote_node);
+    promote_node->eval_jacobian(promote_node);
+
+    /* Jacobian is 4x2:
+     * Row 0: [1, 0] (output 0 depends on input 0)
+     * Row 1: [0, 1] (output 1 depends on input 1)
+     * Row 2: [1, 0] (output 2 depends on input 0)
+     * Row 3: [0, 1] (output 3 depends on input 1)
+     */
+    double expected_x[4] = {1.0, 1.0, 1.0, 1.0};
+    int expected_p[5] = {0, 1, 2, 3, 4};
+    int expected_i[4] = {0, 1, 0, 1};
+
+    mu_assert("promote vector jacobian vals fail",
+              cmp_double_array(promote_node->jacobian->x, expected_x, 4));
+    mu_assert("promote vector jacobian rows fail",
+              cmp_int_array(promote_node->jacobian->p, expected_p, 5));
+    mu_assert("promote vector jacobian cols fail",
+              cmp_int_array(promote_node->jacobian->i, expected_i, 4));
+
+    free_expr(promote_node);
+    free_expr(var);
+    return 0;
+}

From 279c582d006e1f40226c19435a213e70c1bf10a8 Mon Sep 17 00:00:00 2001
From: William Zijie Zhang <william@gridmatic.com>
Date: Mon, 5 Jan 2026 21:19:37 -0500
Subject: [PATCH 09/27] cleanup a bit convert.py

---
 python/convert.py | 35 ++++++++++++++---------------------
 1 file changed, 14 insertions(+), 21 deletions(-)

diff --git a/python/convert.py b/python/convert.py
index a7a6166..3d7a1de 100644
--- a/python/convert.py
+++ b/python/convert.py
@@ -14,16 +14,25 @@ def _chain_add(children):
 
 
 # Mapping from CVXPY atom names to C diff engine functions
+# Converters receive (expr, children) where expr is the CVXPY expression
 ATOM_CONVERTERS = {
     # Elementwise unary
-    "log": lambda child: diffengine.make_log(child),
-    "exp": lambda child: diffengine.make_exp(child),
+    "log": lambda expr, children: diffengine.make_log(children[0]),
+    "exp": lambda expr, children: diffengine.make_exp(children[0]),
+
+    # Affine unary
+    "NegExpression": lambda expr, children: diffengine.make_neg(children[0]),
+    "Promote": lambda expr, children: diffengine.make_promote(
+        children[0],
+        expr.shape[0] if len(expr.shape) >= 1 else 1,
+        expr.shape[1] if len(expr.shape) >= 2 else 1,
+    ),
 
     # N-ary (handles 2+ args)
-    "AddExpression": _chain_add,
+    "AddExpression": lambda expr, children: _chain_add(children),
 
     # Reductions
-    "Sum": lambda child: diffengine.make_sum(child, -1),  # axis=-1 = sum all
+    "Sum": lambda expr, children: diffengine.make_sum(children[0], -1),
 }
 
 
@@ -72,25 +81,9 @@ def _convert_expr(expr, var_dict: dict, n_vars: int):
     # Recursive case: atoms
     atom_name = type(expr).__name__
 
-    # Handle NegExpression using neg atom
-    if atom_name == "NegExpression":
-        child = _convert_expr(expr.args[0], var_dict, n_vars)
-        return diffengine.make_neg(child)
-
-    # Handle Promote (broadcasts scalar/vector to larger shape)
-    if atom_name == "Promote":
-        child = _convert_expr(expr.args[0], var_dict, n_vars)
-        d1 = expr.shape[0] if len(expr.shape) >= 1 else 1
-        d2 = expr.shape[1] if len(expr.shape) >= 2 else 1
-        return diffengine.make_promote(child, d1, d2)
-
     if atom_name in ATOM_CONVERTERS:
         children = [_convert_expr(arg, var_dict, n_vars) for arg in expr.args]
-        converter = ATOM_CONVERTERS[atom_name]
-        # N-ary ops (like AddExpression) take list, unary ops take single arg
-        if atom_name == "AddExpression":
-            return converter(children)
-        return converter(*children) if len(children) > 1 else converter(children[0])
+        return ATOM_CONVERTERS[atom_name](expr, children)
 
     raise NotImplementedError(f"Atom '{atom_name}' not supported")
 

From 2721ae86454238171da3a20e9bd9b795c6370452 Mon Sep 17 00:00:00 2001
From: William Zijie Zhang <william@gridmatic.com>
Date: Tue, 6 Jan 2026 12:40:27 -0500
Subject: [PATCH 10/27] remove dead code in problem.c

---
 src/problem.c | 56 ---------------------------------------------------
 1 file changed, 56 deletions(-)

diff --git a/src/problem.c b/src/problem.c
index 5fc9ad1..bb27567 100644
--- a/src/problem.c
+++ b/src/problem.c
@@ -2,62 +2,6 @@
 #include <stdlib.h>
 #include <string.h>
 
-/* Simple visited set for tracking freed nodes (handles up to 1024 unique nodes) */
-#define MAX_VISITED 1024
-
-typedef struct
-{
-    expr *nodes[MAX_VISITED];
-    int count;
-} VisitedSet;
-
-static void visited_init(VisitedSet *v)
-{
-    v->count = 0;
-}
-
-static int visited_contains(VisitedSet *v, expr *node)
-{
-    for (int i = 0; i < v->count; i++)
-    {
-        if (v->nodes[i] == node) return 1;
-    }
-    return 0;
-}
-
-static void visited_add(VisitedSet *v, expr *node)
-{
-    if (v->count < MAX_VISITED)
-    {
-        v->nodes[v->count++] = node;
-    }
-}
-
-/* Release refs and free nodes, tracking visited to handle sharing */
-static void free_expr_tree_visited(expr *node, VisitedSet *visited)
-{
-    if (node == NULL || visited_contains(visited, node)) return;
-    visited_add(visited, node);
-
-    /* Recursively process children first */
-    free_expr_tree_visited(node->left, visited);
-    free_expr_tree_visited(node->right, visited);
-
-    /* Free this node's resources */
-    free(node->value);
-    free_csr_matrix(node->jacobian);
-    free_csr_matrix(node->wsum_hess);
-    free(node->dwork);
-    free(node->iwork);
-
-    if (node->free_type_data)
-    {
-        node->free_type_data(node);
-    }
-
-    free(node);
-}
-
 problem *new_problem(expr *objective, expr **constraints, int n_constraints)
 {
     problem *prob = (problem *) calloc(1, sizeof(problem));

From 4f0fb2155e6e50185c7cf1f5cee0fba624220d3d Mon Sep 17 00:00:00 2001
From: William Zijie Zhang <william@gridmatic.com>
Date: Tue, 6 Jan 2026 12:57:35 -0500
Subject: [PATCH 11/27] change problem_allocate to init_deriv

---
 include/problem.h                   |  4 ++--
 python/bindings.c                   | 17 ++++-------------
 python/convert.py                   |  6 +++---
 python/tests/test_problem_native.py | 22 ++++++++++++++--------
 src/problem.c                       |  8 +++++---
 tests/problem/test_problem.h        | 13 ++++++++-----
 6 files changed, 36 insertions(+), 34 deletions(-)

diff --git a/include/problem.h b/include/problem.h
index 3dbfc9b..d86bf9a 100644
--- a/include/problem.h
+++ b/include/problem.h
@@ -12,7 +12,7 @@ typedef struct problem
     int n_vars;
     int total_constraint_size;
 
-    /* Allocated by problem_allocate */
+    /* Allocated by problem_init_derivatives */
     double *constraint_values;
     double *gradient_values;
     CSR_Matrix *stacked_jac;
@@ -20,7 +20,7 @@ typedef struct problem
 
 /* Retains objective and constraints (shared ownership with caller) */
 problem *new_problem(expr *objective, expr **constraints, int n_constraints);
-void problem_allocate(problem *prob, const double *u);
+void problem_init_derivatives(problem *prob);
 void free_problem(problem *prob);
 
 double problem_objective_forward(problem *prob, const double *u);
diff --git a/python/bindings.c b/python/bindings.c
index 334e270..f5f4574 100644
--- a/python/bindings.c
+++ b/python/bindings.c
@@ -437,11 +437,10 @@ static PyObject *py_make_problem(PyObject *self, PyObject *args)
     return PyCapsule_New(prob, PROBLEM_CAPSULE_NAME, problem_capsule_destructor);
 }
 
-static PyObject *py_problem_allocate(PyObject *self, PyObject *args)
+static PyObject *py_problem_init_derivatives(PyObject *self, PyObject *args)
 {
     PyObject *prob_capsule;
-    PyObject *u_obj;
-    if (!PyArg_ParseTuple(args, "OO", &prob_capsule, &u_obj))
+    if (!PyArg_ParseTuple(args, "O", &prob_capsule))
     {
         return NULL;
     }
@@ -454,16 +453,8 @@ static PyObject *py_problem_allocate(PyObject *self, PyObject *args)
         return NULL;
     }
 
-    PyArrayObject *u_array =
-        (PyArrayObject *) PyArray_FROM_OTF(u_obj, NPY_DOUBLE, NPY_ARRAY_IN_ARRAY);
-    if (!u_array)
-    {
-        return NULL;
-    }
-
-    problem_allocate(prob, (const double *) PyArray_DATA(u_array));
+    problem_init_derivatives(prob);
 
-    Py_DECREF(u_array);
     Py_RETURN_NONE;
 }
 
@@ -662,7 +653,7 @@ static PyMethodDef DNLPMethods[] = {
     {"jacobian", py_jacobian, METH_VARARGS,
      "Compute jacobian and return CSR components"},
     {"make_problem", py_make_problem, METH_VARARGS, "Create problem from objective and constraints"},
-    {"problem_allocate", py_problem_allocate, METH_VARARGS, "Allocate problem resources"},
+    {"problem_init_derivatives", py_problem_init_derivatives, METH_VARARGS, "Initialize derivative structures"},
     {"problem_objective_forward", py_problem_objective_forward, METH_VARARGS, "Evaluate objective only"},
     {"problem_constraint_forward", py_problem_constraint_forward, METH_VARARGS, "Evaluate constraints only"},
     {"problem_gradient", py_problem_gradient, METH_VARARGS, "Compute objective gradient"},
diff --git a/python/convert.py b/python/convert.py
index 3d7a1de..4c344ed 100644
--- a/python/convert.py
+++ b/python/convert.py
@@ -138,9 +138,9 @@ def __init__(self, cvxpy_problem: cp.Problem):
         self._capsule = diffengine.make_problem(c_obj, c_constraints)
         self._allocated = False
 
-    def allocate(self, u: np.ndarray):
-        """Allocate internal buffers. Must be called before forward/gradient/jacobian."""
-        diffengine.problem_allocate(self._capsule, u)
+    def init_derivatives(self):
+        """Initialize derivative structures. Must be called before gradient/jacobian."""
+        diffengine.problem_init_derivatives(self._capsule)
         self._allocated = True
 
     def objective_forward(self, u: np.ndarray) -> float:
diff --git a/python/tests/test_problem_native.py b/python/tests/test_problem_native.py
index a7ce1c4..56d9842 100644
--- a/python/tests/test_problem_native.py
+++ b/python/tests/test_problem_native.py
@@ -13,7 +13,7 @@ def test_problem_objective_forward():
 
     prob = diffengine.make_problem(objective, constraints)
     u = np.array([1.0, 2.0, 3.0])
-    diffengine.problem_allocate(prob, u)
+    diffengine.problem_init_derivatives(prob)
 
     obj_val = diffengine.problem_objective_forward(prob, u)
     constraint_vals = diffengine.problem_constraint_forward(prob, u)
@@ -37,7 +37,7 @@ def test_problem_constraint_forward():
 
     prob = diffengine.make_problem(objective, constraints)
     u = np.array([2.0, 4.0])
-    diffengine.problem_allocate(prob, u)
+    diffengine.problem_init_derivatives(prob)
 
     constraint_vals = diffengine.problem_constraint_forward(prob, u)
 
@@ -55,8 +55,9 @@ def test_problem_gradient():
 
     prob = diffengine.make_problem(objective, [])
     u = np.array([1.0, 2.0, 4.0])
-    diffengine.problem_allocate(prob, u)
+    diffengine.problem_init_derivatives(prob)
 
+    diffengine.problem_objective_forward(prob, u)
     grad = diffengine.problem_gradient(prob, u)
     expected_grad = 1.0 / u
     assert np.allclose(grad, expected_grad)
@@ -72,8 +73,9 @@ def test_problem_jacobian():
 
     prob = diffengine.make_problem(objective, constraints)
     u = np.array([2.0, 4.0])
-    diffengine.problem_allocate(prob, u)
+    diffengine.problem_init_derivatives(prob)
 
+    diffengine.problem_constraint_forward(prob, u)
     data, indices, indptr, shape = diffengine.problem_jacobian(prob, u)
     jac = sparse.csr_matrix((data, indices, indptr), shape=shape)
 
@@ -90,16 +92,18 @@ def test_problem_no_constraints():
 
     prob = diffengine.make_problem(objective, [])
     u = np.array([1.0, 2.0, 3.0])
-    diffengine.problem_allocate(prob, u)
+    diffengine.problem_init_derivatives(prob)
 
     obj_val = diffengine.problem_objective_forward(prob, u)
     constraint_vals = diffengine.problem_constraint_forward(prob, u)
     assert np.allclose(obj_val, np.sum(np.log(u)))
     assert len(constraint_vals) == 0
 
+    diffengine.problem_objective_forward(prob, u)
     grad = diffengine.problem_gradient(prob, u)
     assert np.allclose(grad, 1.0 / u)
 
+    diffengine.problem_constraint_forward(prob, u)
     data, indices, indptr, shape = diffengine.problem_jacobian(prob, u)
     jac = sparse.csr_matrix((data, indices, indptr), shape=shape)
     assert jac.shape == (0, 3)
@@ -122,7 +126,7 @@ def test_problem_multiple_constraints():
     
     prob = diffengine.make_problem(objective, constraints)
     u = np.array([1.0, 2.0, 3.0])
-    diffengine.problem_allocate(prob, u)
+    diffengine.problem_init_derivatives(prob)
     
     # Test forward pass
     obj_val = diffengine.problem_objective_forward(prob, u)
@@ -131,13 +135,15 @@ def test_problem_multiple_constraints():
     expected_constraints = np.concatenate([np.log(u), np.exp(u)])
     assert np.allclose(obj_val, expected_obj)
     assert np.allclose(constraint_vals, expected_constraints)
-    
+
     # Test gradient
+    diffengine.problem_objective_forward(prob, u)
     grad = diffengine.problem_gradient(prob, u)
     expected_grad = 1.0 / u
     assert np.allclose(grad, expected_grad)
-    
+
     # Test Jacobian
+    diffengine.problem_constraint_forward(prob, u)
     data, indices, indptr, shape = diffengine.problem_jacobian(prob, u)
     jac = sparse.csr_matrix((data, indices, indptr), shape=shape)
     
diff --git a/src/problem.c b/src/problem.c
index bb27567..63281b5 100644
--- a/src/problem.c
+++ b/src/problem.c
@@ -44,7 +44,7 @@ problem *new_problem(expr *objective, expr **constraints, int n_constraints)
     return prob;
 }
 
-void problem_allocate(problem *prob, const double *u)
+void problem_init_derivatives(problem *prob)
 {
     /* 1. Allocate constraint values array */
     if (prob->total_constraint_size > 0)
@@ -56,7 +56,6 @@ void problem_allocate(problem *prob, const double *u)
     prob->gradient_values = (double *) calloc(prob->n_vars, sizeof(double));
 
     /* 3. Initialize objective jacobian */
-    prob->objective->forward(prob->objective, u);
     prob->objective->jacobian_init(prob->objective);
 
     /* 4. Initialize constraint jacobians and count total nnz */
@@ -64,7 +63,6 @@ void problem_allocate(problem *prob, const double *u)
     for (int i = 0; i < prob->n_constraints; i++)
     {
         expr *c = prob->constraints[i];
-        c->forward(c, u);
         c->jacobian_init(c);
         total_nnz += c->jacobian->nnz;
     }
@@ -74,6 +72,10 @@ void problem_allocate(problem *prob, const double *u)
     {
         prob->stacked_jac = new_csr_matrix(prob->total_constraint_size, prob->n_vars, total_nnz);
     }
+
+    /* TODO: 6. Initialize objective wsum_hess */
+
+    /* TODO: 7. Initialize constraint wsum_hess */
 }
 
 void free_problem(problem *prob)
diff --git a/tests/problem/test_problem.h b/tests/problem/test_problem.h
index 814a3c6..b5e1c90 100644
--- a/tests/problem/test_problem.h
+++ b/tests/problem/test_problem.h
@@ -63,7 +63,7 @@ const char *test_problem_objective_forward(void)
     problem *prob = new_problem(objective, constraints, 1);
 
     double u[3] = {1.0, 2.0, 3.0};
-    problem_allocate(prob, u);
+    problem_init_derivatives(prob);
 
     double obj_val = problem_objective_forward(prob, u);
 
@@ -115,7 +115,7 @@ const char *test_problem_constraint_forward(void)
     problem *prob = new_problem(objective, constraints, 2);
 
     double u[2] = {2.0, 4.0};
-    problem_allocate(prob, u);
+    problem_init_derivatives(prob);
 
     double *constraint_vals = problem_constraint_forward(prob, u);
 
@@ -149,8 +149,9 @@ const char *test_problem_gradient(void)
     problem *prob = new_problem(objective, NULL, 0);
 
     double u[3] = {1.0, 2.0, 4.0};
-    problem_allocate(prob, u);
+    problem_init_derivatives(prob);
 
+    problem_objective_forward(prob, u);
     double *grad = problem_gradient(prob, u);
 
     /* Expected gradient: [1/1, 1/2, 1/4] = [1.0, 0.5, 0.25] */
@@ -187,8 +188,9 @@ const char *test_problem_jacobian(void)
     problem *prob = new_problem(objective, constraints, 1);
 
     double u[2] = {2.0, 4.0};
-    problem_allocate(prob, u);
+    problem_init_derivatives(prob);
 
+    problem_constraint_forward(prob, u);
     CSR_Matrix *jac = problem_jacobian(prob, u);
 
     /* Check dimensions */
@@ -254,8 +256,9 @@ const char *test_problem_jacobian_multi(void)
     problem *prob = new_problem(objective, constraints, 2);
 
     double u[2] = {2.0, 4.0};
-    problem_allocate(prob, u);
+    problem_init_derivatives(prob);
 
+    problem_constraint_forward(prob, u);
     CSR_Matrix *jac = problem_jacobian(prob, u);
 
     /* Check dimensions: 4 rows (2 + 2), 2 cols */

From 5b27f2548a687b2fbc3d8d8db91bc7178b9721c0 Mon Sep 17 00:00:00 2001
From: William Zijie Zhang <william@gridmatic.com>
Date: Tue, 6 Jan 2026 13:08:15 -0500
Subject: [PATCH 12/27] refactor python bindings into atoms/ and problem/
 folders
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Split Python binding functions into separate header files:
- atoms/: one file per atom type (variable, constant, linear, log,
  exp, add, sum, neg, promote)
- problem/: one file per method (make_problem, init_derivatives,
  objective_forward, constraint_forward, gradient, jacobian)

Removed standalone forward/jacobian bindings (use problem_* instead).
bindings.c now only contains includes and module definition.

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
---
 python/atoms/add.h                  |  30 ++
 python/atoms/common.h               |  23 +
 python/atoms/constant.h             |  34 ++
 python/atoms/exp.h                  |  29 ++
 python/atoms/linear.h               |  60 +++
 python/atoms/log.h                  |  29 ++
 python/atoms/neg.h                  |  29 ++
 python/atoms/promote.h              |  30 ++
 python/atoms/sum.h                  |  30 ++
 python/atoms/variable.h             |  23 +
 python/bindings.c                   | 647 +---------------------------
 python/problem/common.h             |  24 ++
 python/problem/constraint_forward.h |  56 +++
 python/problem/gradient.h           |  45 ++
 python/problem/init_derivatives.h   |  27 ++
 python/problem/jacobian.h           |  68 +++
 python/problem/make_problem.h       |  64 +++
 python/problem/objective_forward.h  |  36 ++
 18 files changed, 655 insertions(+), 629 deletions(-)
 create mode 100644 python/atoms/add.h
 create mode 100644 python/atoms/common.h
 create mode 100644 python/atoms/constant.h
 create mode 100644 python/atoms/exp.h
 create mode 100644 python/atoms/linear.h
 create mode 100644 python/atoms/log.h
 create mode 100644 python/atoms/neg.h
 create mode 100644 python/atoms/promote.h
 create mode 100644 python/atoms/sum.h
 create mode 100644 python/atoms/variable.h
 create mode 100644 python/problem/common.h
 create mode 100644 python/problem/constraint_forward.h
 create mode 100644 python/problem/gradient.h
 create mode 100644 python/problem/init_derivatives.h
 create mode 100644 python/problem/jacobian.h
 create mode 100644 python/problem/make_problem.h
 create mode 100644 python/problem/objective_forward.h

diff --git a/python/atoms/add.h b/python/atoms/add.h
new file mode 100644
index 0000000..4b86b79
--- /dev/null
+++ b/python/atoms/add.h
@@ -0,0 +1,30 @@
+#ifndef ATOM_ADD_H
+#define ATOM_ADD_H
+
+#include "common.h"
+
+static PyObject *py_make_add(PyObject *self, PyObject *args)
+{
+    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_add(left, right);
+    if (!node)
+    {
+        PyErr_SetString(PyExc_RuntimeError, "failed to create add node");
+        return NULL;
+    }
+    return PyCapsule_New(node, EXPR_CAPSULE_NAME, expr_capsule_destructor);
+}
+
+#endif /* ATOM_ADD_H */
diff --git a/python/atoms/common.h b/python/atoms/common.h
new file mode 100644
index 0000000..83101dd
--- /dev/null
+++ b/python/atoms/common.h
@@ -0,0 +1,23 @@
+#ifndef ATOMS_COMMON_H
+#define ATOMS_COMMON_H
+
+#define PY_SSIZE_T_CLEAN
+#include <Python.h>
+#include <numpy/arrayobject.h>
+
+#include "affine.h"
+#include "elementwise_univariate.h"
+#include "expr.h"
+
+#define EXPR_CAPSULE_NAME "DNLP_EXPR"
+
+static void expr_capsule_destructor(PyObject *capsule)
+{
+    expr *node = (expr *) PyCapsule_GetPointer(capsule, EXPR_CAPSULE_NAME);
+    if (node)
+    {
+        free_expr(node);
+    }
+}
+
+#endif /* ATOMS_COMMON_H */
diff --git a/python/atoms/constant.h b/python/atoms/constant.h
new file mode 100644
index 0000000..8062647
--- /dev/null
+++ b/python/atoms/constant.h
@@ -0,0 +1,34 @@
+#ifndef ATOM_CONSTANT_H
+#define ATOM_CONSTANT_H
+
+#include "common.h"
+
+static PyObject *py_make_constant(PyObject *self, PyObject *args)
+{
+    int d1, d2, n_vars;
+    PyObject *values_obj;
+    if (!PyArg_ParseTuple(args, "iiiO", &d1, &d2, &n_vars, &values_obj))
+    {
+        return NULL;
+    }
+
+    PyArrayObject *values_array =
+        (PyArrayObject *) PyArray_FROM_OTF(values_obj, NPY_DOUBLE, NPY_ARRAY_IN_ARRAY);
+    if (!values_array)
+    {
+        return NULL;
+    }
+
+    expr *node =
+        new_constant(d1, d2, n_vars, (const double *) PyArray_DATA(values_array));
+    Py_DECREF(values_array);
+
+    if (!node)
+    {
+        PyErr_SetString(PyExc_RuntimeError, "failed to create constant node");
+        return NULL;
+    }
+    return PyCapsule_New(node, EXPR_CAPSULE_NAME, expr_capsule_destructor);
+}
+
+#endif /* ATOM_CONSTANT_H */
diff --git a/python/atoms/exp.h b/python/atoms/exp.h
new file mode 100644
index 0000000..9505daf
--- /dev/null
+++ b/python/atoms/exp.h
@@ -0,0 +1,29 @@
+#ifndef ATOM_EXP_H
+#define ATOM_EXP_H
+
+#include "common.h"
+
+static PyObject *py_make_exp(PyObject *self, PyObject *args)
+{
+    PyObject *child_capsule;
+    if (!PyArg_ParseTuple(args, "O", &child_capsule))
+    {
+        return NULL;
+    }
+    expr *child = (expr *) PyCapsule_GetPointer(child_capsule, EXPR_CAPSULE_NAME);
+    if (!child)
+    {
+        PyErr_SetString(PyExc_ValueError, "invalid child capsule");
+        return NULL;
+    }
+
+    expr *node = new_exp(child);
+    if (!node)
+    {
+        PyErr_SetString(PyExc_RuntimeError, "failed to create exp node");
+        return NULL;
+    }
+    return PyCapsule_New(node, EXPR_CAPSULE_NAME, expr_capsule_destructor);
+}
+
+#endif /* ATOM_EXP_H */
diff --git a/python/atoms/linear.h b/python/atoms/linear.h
new file mode 100644
index 0000000..269b34f
--- /dev/null
+++ b/python/atoms/linear.h
@@ -0,0 +1,60 @@
+#ifndef ATOM_LINEAR_H
+#define ATOM_LINEAR_H
+
+#include "common.h"
+
+static PyObject *py_make_linear(PyObject *self, PyObject *args)
+{
+    PyObject *child_capsule;
+    PyObject *data_obj, *indices_obj, *indptr_obj;
+    int m, n;
+    if (!PyArg_ParseTuple(args, "OOOOii", &child_capsule, &data_obj, &indices_obj,
+                          &indptr_obj, &m, &n))
+    {
+        return NULL;
+    }
+
+    expr *child = (expr *) PyCapsule_GetPointer(child_capsule, EXPR_CAPSULE_NAME);
+    if (!child)
+    {
+        PyErr_SetString(PyExc_ValueError, "invalid child capsule");
+        return NULL;
+    }
+
+    PyArrayObject *data_array =
+        (PyArrayObject *) PyArray_FROM_OTF(data_obj, NPY_DOUBLE, NPY_ARRAY_IN_ARRAY);
+    PyArrayObject *indices_array =
+        (PyArrayObject *) PyArray_FROM_OTF(indices_obj, NPY_INT32, NPY_ARRAY_IN_ARRAY);
+    PyArrayObject *indptr_array =
+        (PyArrayObject *) PyArray_FROM_OTF(indptr_obj, NPY_INT32, NPY_ARRAY_IN_ARRAY);
+
+    if (!data_array || !indices_array || !indptr_array)
+    {
+        Py_XDECREF(data_array);
+        Py_XDECREF(indices_array);
+        Py_XDECREF(indptr_array);
+        return NULL;
+    }
+
+    int nnz = (int) PyArray_SIZE(data_array);
+    CSR_Matrix *A = new_csr_matrix(m, n, nnz);
+    memcpy(A->x, PyArray_DATA(data_array), nnz * sizeof(double));
+    memcpy(A->i, PyArray_DATA(indices_array), nnz * sizeof(int));
+    memcpy(A->p, PyArray_DATA(indptr_array), (m + 1) * sizeof(int));
+
+    Py_DECREF(data_array);
+    Py_DECREF(indices_array);
+    Py_DECREF(indptr_array);
+
+    expr *node = new_linear(child, A);
+    free_csr_matrix(A);
+
+    if (!node)
+    {
+        PyErr_SetString(PyExc_RuntimeError, "failed to create linear node");
+        return NULL;
+    }
+    return PyCapsule_New(node, EXPR_CAPSULE_NAME, expr_capsule_destructor);
+}
+
+#endif /* ATOM_LINEAR_H */
diff --git a/python/atoms/log.h b/python/atoms/log.h
new file mode 100644
index 0000000..3d0314e
--- /dev/null
+++ b/python/atoms/log.h
@@ -0,0 +1,29 @@
+#ifndef ATOM_LOG_H
+#define ATOM_LOG_H
+
+#include "common.h"
+
+static PyObject *py_make_log(PyObject *self, PyObject *args)
+{
+    PyObject *child_capsule;
+    if (!PyArg_ParseTuple(args, "O", &child_capsule))
+    {
+        return NULL;
+    }
+    expr *child = (expr *) PyCapsule_GetPointer(child_capsule, EXPR_CAPSULE_NAME);
+    if (!child)
+    {
+        PyErr_SetString(PyExc_ValueError, "invalid child capsule");
+        return NULL;
+    }
+
+    expr *node = new_log(child);
+    if (!node)
+    {
+        PyErr_SetString(PyExc_RuntimeError, "failed to create log node");
+        return NULL;
+    }
+    return PyCapsule_New(node, EXPR_CAPSULE_NAME, expr_capsule_destructor);
+}
+
+#endif /* ATOM_LOG_H */
diff --git a/python/atoms/neg.h b/python/atoms/neg.h
new file mode 100644
index 0000000..bda69d0
--- /dev/null
+++ b/python/atoms/neg.h
@@ -0,0 +1,29 @@
+#ifndef ATOM_NEG_H
+#define ATOM_NEG_H
+
+#include "common.h"
+
+static PyObject *py_make_neg(PyObject *self, PyObject *args)
+{
+    PyObject *child_capsule;
+    if (!PyArg_ParseTuple(args, "O", &child_capsule))
+    {
+        return NULL;
+    }
+    expr *child = (expr *) PyCapsule_GetPointer(child_capsule, EXPR_CAPSULE_NAME);
+    if (!child)
+    {
+        PyErr_SetString(PyExc_ValueError, "invalid child capsule");
+        return NULL;
+    }
+
+    expr *node = new_neg(child);
+    if (!node)
+    {
+        PyErr_SetString(PyExc_RuntimeError, "failed to create neg node");
+        return NULL;
+    }
+    return PyCapsule_New(node, EXPR_CAPSULE_NAME, expr_capsule_destructor);
+}
+
+#endif /* ATOM_NEG_H */
diff --git a/python/atoms/promote.h b/python/atoms/promote.h
new file mode 100644
index 0000000..fb258b4
--- /dev/null
+++ b/python/atoms/promote.h
@@ -0,0 +1,30 @@
+#ifndef ATOM_PROMOTE_H
+#define ATOM_PROMOTE_H
+
+#include "common.h"
+
+static PyObject *py_make_promote(PyObject *self, PyObject *args)
+{
+    PyObject *child_capsule;
+    int d1, d2;
+    if (!PyArg_ParseTuple(args, "Oii", &child_capsule, &d1, &d2))
+    {
+        return NULL;
+    }
+    expr *child = (expr *) PyCapsule_GetPointer(child_capsule, EXPR_CAPSULE_NAME);
+    if (!child)
+    {
+        PyErr_SetString(PyExc_ValueError, "invalid child capsule");
+        return NULL;
+    }
+
+    expr *node = new_promote(child, d1, d2);
+    if (!node)
+    {
+        PyErr_SetString(PyExc_RuntimeError, "failed to create promote node");
+        return NULL;
+    }
+    return PyCapsule_New(node, EXPR_CAPSULE_NAME, expr_capsule_destructor);
+}
+
+#endif /* ATOM_PROMOTE_H */
diff --git a/python/atoms/sum.h b/python/atoms/sum.h
new file mode 100644
index 0000000..8a8480e
--- /dev/null
+++ b/python/atoms/sum.h
@@ -0,0 +1,30 @@
+#ifndef ATOM_SUM_H
+#define ATOM_SUM_H
+
+#include "common.h"
+
+static PyObject *py_make_sum(PyObject *self, PyObject *args)
+{
+    PyObject *child_capsule;
+    int axis;
+    if (!PyArg_ParseTuple(args, "Oi", &child_capsule, &axis))
+    {
+        return NULL;
+    }
+    expr *child = (expr *) PyCapsule_GetPointer(child_capsule, EXPR_CAPSULE_NAME);
+    if (!child)
+    {
+        PyErr_SetString(PyExc_ValueError, "invalid child capsule");
+        return NULL;
+    }
+
+    expr *node = new_sum(child, axis);
+    if (!node)
+    {
+        PyErr_SetString(PyExc_RuntimeError, "failed to create sum node");
+        return NULL;
+    }
+    return PyCapsule_New(node, EXPR_CAPSULE_NAME, expr_capsule_destructor);
+}
+
+#endif /* ATOM_SUM_H */
diff --git a/python/atoms/variable.h b/python/atoms/variable.h
new file mode 100644
index 0000000..5438c44
--- /dev/null
+++ b/python/atoms/variable.h
@@ -0,0 +1,23 @@
+#ifndef ATOM_VARIABLE_H
+#define ATOM_VARIABLE_H
+
+#include "common.h"
+
+static PyObject *py_make_variable(PyObject *self, PyObject *args)
+{
+    int d1, d2, var_id, n_vars;
+    if (!PyArg_ParseTuple(args, "iiii", &d1, &d2, &var_id, &n_vars))
+    {
+        return NULL;
+    }
+
+    expr *node = new_variable(d1, d2, var_id, n_vars);
+    if (!node)
+    {
+        PyErr_SetString(PyExc_RuntimeError, "failed to create variable node");
+        return NULL;
+    }
+    return PyCapsule_New(node, EXPR_CAPSULE_NAME, expr_capsule_destructor);
+}
+
+#endif /* ATOM_VARIABLE_H */
diff --git a/python/bindings.c b/python/bindings.c
index f5f4574..e153e0e 100644
--- a/python/bindings.c
+++ b/python/bindings.c
@@ -1,16 +1,25 @@
 #define PY_SSIZE_T_CLEAN
 #include <Python.h>
 #include <numpy/arrayobject.h>
-#include <stdbool.h>
 
-#include "affine.h"
-#include "elementwise_univariate.h"
-#include "expr.h"
-#include "problem.h"
-
-// Capsule name for expr* pointers
-#define EXPR_CAPSULE_NAME "DNLP_EXPR"
-#define PROBLEM_CAPSULE_NAME "DNLP_PROBLEM"
+/* Include atom bindings */
+#include "atoms/variable.h"
+#include "atoms/constant.h"
+#include "atoms/linear.h"
+#include "atoms/log.h"
+#include "atoms/exp.h"
+#include "atoms/add.h"
+#include "atoms/sum.h"
+#include "atoms/neg.h"
+#include "atoms/promote.h"
+
+/* Include problem bindings */
+#include "problem/make_problem.h"
+#include "problem/init_derivatives.h"
+#include "problem/objective_forward.h"
+#include "problem/constraint_forward.h"
+#include "problem/gradient.h"
+#include "problem/jacobian.h"
 
 static int numpy_initialized = 0;
 
@@ -22,623 +31,6 @@ static int ensure_numpy(void)
     return 0;
 }
 
-static void expr_capsule_destructor(PyObject *capsule)
-{
-    expr *node = (expr *) PyCapsule_GetPointer(capsule, EXPR_CAPSULE_NAME);
-    if (node)
-    {
-        free_expr(node);
-    }
-}
-
-static PyObject *py_make_variable(PyObject *self, PyObject *args)
-{
-    int d1, d2, var_id, n_vars;
-    if (!PyArg_ParseTuple(args, "iiii", &d1, &d2, &var_id, &n_vars))
-    {
-        return NULL;
-    }
-
-    expr *node = new_variable(d1, d2, var_id, n_vars);
-    if (!node)
-    {
-        PyErr_SetString(PyExc_RuntimeError, "failed to create variable node");
-        return NULL;
-    }
-    return PyCapsule_New(node, EXPR_CAPSULE_NAME, expr_capsule_destructor);
-}
-
-static PyObject *py_make_constant(PyObject *self, PyObject *args)
-{
-    int d1, d2, n_vars;
-    PyObject *values_obj;
-    if (!PyArg_ParseTuple(args, "iiiO", &d1, &d2, &n_vars, &values_obj))
-    {
-        return NULL;
-    }
-
-    PyArrayObject *values_array =
-        (PyArrayObject *) PyArray_FROM_OTF(values_obj, NPY_DOUBLE, NPY_ARRAY_IN_ARRAY);
-    if (!values_array)
-    {
-        return NULL;
-    }
-
-    expr *node =
-        new_constant(d1, d2, n_vars, (const double *) PyArray_DATA(values_array));
-    Py_DECREF(values_array);
-
-    if (!node)
-    {
-        PyErr_SetString(PyExc_RuntimeError, "failed to create constant node");
-        return NULL;
-    }
-    return PyCapsule_New(node, EXPR_CAPSULE_NAME, expr_capsule_destructor);
-}
-
-static PyObject *py_make_linear(PyObject *self, PyObject *args)
-{
-    PyObject *child_capsule;
-    PyObject *data_obj, *indices_obj, *indptr_obj;
-    int m, n;
-    if (!PyArg_ParseTuple(args, "OOOOii", &child_capsule, &data_obj, &indices_obj,
-                          &indptr_obj, &m, &n))
-    {
-        return NULL;
-    }
-
-    expr *child = (expr *) PyCapsule_GetPointer(child_capsule, EXPR_CAPSULE_NAME);
-    if (!child)
-    {
-        PyErr_SetString(PyExc_ValueError, "invalid child capsule");
-        return NULL;
-    }
-
-    PyArrayObject *data_array =
-        (PyArrayObject *) PyArray_FROM_OTF(data_obj, NPY_DOUBLE, NPY_ARRAY_IN_ARRAY);
-    PyArrayObject *indices_array =
-        (PyArrayObject *) PyArray_FROM_OTF(indices_obj, NPY_INT32, NPY_ARRAY_IN_ARRAY);
-    PyArrayObject *indptr_array =
-        (PyArrayObject *) PyArray_FROM_OTF(indptr_obj, NPY_INT32, NPY_ARRAY_IN_ARRAY);
-
-    if (!data_array || !indices_array || !indptr_array)
-    {
-        Py_XDECREF(data_array);
-        Py_XDECREF(indices_array);
-        Py_XDECREF(indptr_array);
-        return NULL;
-    }
-
-    int nnz = (int) PyArray_SIZE(data_array);
-    CSR_Matrix *A = new_csr_matrix(m, n, nnz);
-    memcpy(A->x, PyArray_DATA(data_array), nnz * sizeof(double));
-    memcpy(A->i, PyArray_DATA(indices_array), nnz * sizeof(int));
-    memcpy(A->p, PyArray_DATA(indptr_array), (m + 1) * sizeof(int));
-
-    Py_DECREF(data_array);
-    Py_DECREF(indices_array);
-    Py_DECREF(indptr_array);
-
-    expr *node = new_linear(child, A);
-    free_csr_matrix(A);
-
-    if (!node)
-    {
-        PyErr_SetString(PyExc_RuntimeError, "failed to create linear node");
-        return NULL;
-    }
-    return PyCapsule_New(node, EXPR_CAPSULE_NAME, expr_capsule_destructor);
-}
-
-static PyObject *py_make_log(PyObject *self, PyObject *args)
-{
-    PyObject *child_capsule;
-    if (!PyArg_ParseTuple(args, "O", &child_capsule))
-    {
-        return NULL;
-    }
-    expr *child = (expr *) PyCapsule_GetPointer(child_capsule, EXPR_CAPSULE_NAME);
-    if (!child)
-    {
-        PyErr_SetString(PyExc_ValueError, "invalid child capsule");
-        return NULL;
-    }
-
-    expr *node = new_log(child);
-    if (!node)
-    {
-        PyErr_SetString(PyExc_RuntimeError, "failed to create log node");
-        return NULL;
-    }
-    return PyCapsule_New(node, EXPR_CAPSULE_NAME, expr_capsule_destructor);
-}
-
-static PyObject *py_make_exp(PyObject *self, PyObject *args)
-{
-    PyObject *child_capsule;
-    if (!PyArg_ParseTuple(args, "O", &child_capsule))
-    {
-        return NULL;
-    }
-    expr *child = (expr *) PyCapsule_GetPointer(child_capsule, EXPR_CAPSULE_NAME);
-    if (!child)
-    {
-        PyErr_SetString(PyExc_ValueError, "invalid child capsule");
-        return NULL;
-    }
-
-    expr *node = new_exp(child);
-    if (!node)
-    {
-        PyErr_SetString(PyExc_RuntimeError, "failed to create exp node");
-        return NULL;
-    }
-    return PyCapsule_New(node, EXPR_CAPSULE_NAME, expr_capsule_destructor);
-}
-
-static PyObject *py_make_add(PyObject *self, PyObject *args)
-{
-    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_add(left, right);
-    if (!node)
-    {
-        PyErr_SetString(PyExc_RuntimeError, "failed to create add node");
-        return NULL;
-    }
-    return PyCapsule_New(node, EXPR_CAPSULE_NAME, expr_capsule_destructor);
-}
-
-static PyObject *py_make_sum(PyObject *self, PyObject *args)
-{
-    PyObject *child_capsule;
-    int axis;
-    if (!PyArg_ParseTuple(args, "Oi", &child_capsule, &axis))
-    {
-        return NULL;
-    }
-    expr *child = (expr *) PyCapsule_GetPointer(child_capsule, EXPR_CAPSULE_NAME);
-    if (!child)
-    {
-        PyErr_SetString(PyExc_ValueError, "invalid child capsule");
-        return NULL;
-    }
-
-    expr *node = new_sum(child, axis);
-    if (!node)
-    {
-        PyErr_SetString(PyExc_RuntimeError, "failed to create sum node");
-        return NULL;
-    }
-    return PyCapsule_New(node, EXPR_CAPSULE_NAME, expr_capsule_destructor);
-}
-
-static PyObject *py_make_neg(PyObject *self, PyObject *args)
-{
-    PyObject *child_capsule;
-    if (!PyArg_ParseTuple(args, "O", &child_capsule))
-    {
-        return NULL;
-    }
-    expr *child = (expr *) PyCapsule_GetPointer(child_capsule, EXPR_CAPSULE_NAME);
-    if (!child)
-    {
-        PyErr_SetString(PyExc_ValueError, "invalid child capsule");
-        return NULL;
-    }
-
-    expr *node = new_neg(child);
-    if (!node)
-    {
-        PyErr_SetString(PyExc_RuntimeError, "failed to create neg node");
-        return NULL;
-    }
-    return PyCapsule_New(node, EXPR_CAPSULE_NAME, expr_capsule_destructor);
-}
-
-static PyObject *py_make_promote(PyObject *self, PyObject *args)
-{
-    PyObject *child_capsule;
-    int d1, d2;
-    if (!PyArg_ParseTuple(args, "Oii", &child_capsule, &d1, &d2))
-    {
-        return NULL;
-    }
-    expr *child = (expr *) PyCapsule_GetPointer(child_capsule, EXPR_CAPSULE_NAME);
-    if (!child)
-    {
-        PyErr_SetString(PyExc_ValueError, "invalid child capsule");
-        return NULL;
-    }
-
-    expr *node = new_promote(child, d1, d2);
-    if (!node)
-    {
-        PyErr_SetString(PyExc_RuntimeError, "failed to create promote node");
-        return NULL;
-    }
-    return PyCapsule_New(node, EXPR_CAPSULE_NAME, expr_capsule_destructor);
-}
-
-static PyObject *py_forward(PyObject *self, PyObject *args)
-{
-    PyObject *node_capsule;
-    PyObject *u_obj;
-    if (!PyArg_ParseTuple(args, "OO", &node_capsule, &u_obj))
-    {
-        return NULL;
-    }
-
-    expr *node = (expr *) PyCapsule_GetPointer(node_capsule, EXPR_CAPSULE_NAME);
-    if (!node)
-    {
-        PyErr_SetString(PyExc_ValueError, "invalid node capsule");
-        return NULL;
-    }
-
-    PyArrayObject *u_array =
-        (PyArrayObject *) PyArray_FROM_OTF(u_obj, NPY_DOUBLE, NPY_ARRAY_IN_ARRAY);
-    if (!u_array)
-    {
-        return NULL;
-    }
-
-    node->forward(node, (const double *) PyArray_DATA(u_array));
-
-    npy_intp size = node->size;
-    PyObject *out = PyArray_SimpleNew(1, &size, NPY_DOUBLE);
-    if (!out)
-    {
-        Py_DECREF(u_array);
-        return NULL;
-    }
-    memcpy(PyArray_DATA((PyArrayObject *) out), node->value, size * sizeof(double));
-
-    Py_DECREF(u_array);
-    return out;
-}
-
-static PyObject *py_jacobian(PyObject *self, PyObject *args)
-{
-    PyObject *node_capsule;
-    PyObject *u_obj;
-    if (!PyArg_ParseTuple(args, "OO", &node_capsule, &u_obj))
-    {
-        return NULL;
-    }
-
-    expr *node = (expr *) PyCapsule_GetPointer(node_capsule, EXPR_CAPSULE_NAME);
-    if (!node)
-    {
-        PyErr_SetString(PyExc_ValueError, "invalid node capsule");
-        return NULL;
-    }
-
-    PyArrayObject *u_array =
-        (PyArrayObject *) PyArray_FROM_OTF(u_obj, NPY_DOUBLE, NPY_ARRAY_IN_ARRAY);
-    if (!u_array)
-    {
-        return NULL;
-    }
-
-    // Run forward pass first (required before jacobian)
-    node->forward(node, (const double *) PyArray_DATA(u_array));
-
-    // Initialize and evaluate jacobian
-    node->jacobian_init(node);
-    node->eval_jacobian(node);
-
-    CSR_Matrix *jac = node->jacobian;
-
-    // Create numpy arrays for CSR components
-    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);
-        Py_DECREF(u_array);
-        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));
-
-    Py_DECREF(u_array);
-
-    // Return tuple: (data, indices, indptr, shape)
-    return Py_BuildValue("(OOO(ii))", data, indices, indptr, jac->m, jac->n);
-}
-
-/* ========== Problem bindings ========== */
-
-static void problem_capsule_destructor(PyObject *capsule)
-{
-    problem *prob = (problem *) PyCapsule_GetPointer(capsule, PROBLEM_CAPSULE_NAME);
-    if (prob)
-    {
-        free_problem(prob);
-    }
-}
-
-static PyObject *py_make_problem(PyObject *self, PyObject *args)
-{
-    PyObject *obj_capsule;
-    PyObject *constraints_list;
-    if (!PyArg_ParseTuple(args, "OO", &obj_capsule, &constraints_list))
-    {
-        return NULL;
-    }
-
-    expr *objective = (expr *) PyCapsule_GetPointer(obj_capsule, EXPR_CAPSULE_NAME);
-    if (!objective)
-    {
-        PyErr_SetString(PyExc_ValueError, "invalid objective capsule");
-        return NULL;
-    }
-
-    if (!PyList_Check(constraints_list))
-    {
-        PyErr_SetString(PyExc_TypeError, "constraints must be a list");
-        return NULL;
-    }
-
-    Py_ssize_t n_constraints = PyList_Size(constraints_list);
-    expr **constraints = NULL;
-    if (n_constraints > 0)
-    {
-        constraints = (expr **) malloc(n_constraints * sizeof(expr *));
-        if (!constraints)
-        {
-            PyErr_NoMemory();
-            return NULL;
-        }
-        for (Py_ssize_t i = 0; i < n_constraints; i++)
-        {
-            PyObject *c_capsule = PyList_GetItem(constraints_list, i);
-            constraints[i] =
-                (expr *) PyCapsule_GetPointer(c_capsule, EXPR_CAPSULE_NAME);
-            if (!constraints[i])
-            {
-                free(constraints);
-                PyErr_SetString(PyExc_ValueError, "invalid constraint capsule");
-                return NULL;
-            }
-        }
-    }
-
-    problem *prob = new_problem(objective, constraints, (int) n_constraints);
-    free(constraints);
-
-    if (!prob)
-    {
-        PyErr_SetString(PyExc_RuntimeError, "failed to create problem");
-        return NULL;
-    }
-
-    return PyCapsule_New(prob, PROBLEM_CAPSULE_NAME, problem_capsule_destructor);
-}
-
-static PyObject *py_problem_init_derivatives(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;
-    }
-
-    problem_init_derivatives(prob);
-
-    Py_RETURN_NONE;
-}
-
-static PyObject *py_problem_objective_forward(PyObject *self, PyObject *args)
-{
-    PyObject *prob_capsule;
-    PyObject *u_obj;
-    if (!PyArg_ParseTuple(args, "OO", &prob_capsule, &u_obj))
-    {
-        return NULL;
-    }
-
-    problem *prob =
-        (problem *) PyCapsule_GetPointer(prob_capsule, PROBLEM_CAPSULE_NAME);
-    if (!prob)
-    {
-        PyErr_SetString(PyExc_ValueError, "invalid problem capsule");
-        return NULL;
-    }
-
-    PyArrayObject *u_array =
-        (PyArrayObject *) PyArray_FROM_OTF(u_obj, NPY_DOUBLE, NPY_ARRAY_IN_ARRAY);
-    if (!u_array)
-    {
-        return NULL;
-    }
-
-    double obj_val = problem_objective_forward(prob, (const double *) PyArray_DATA(u_array));
-
-    Py_DECREF(u_array);
-    return Py_BuildValue("d", obj_val);
-}
-
-static PyObject *py_problem_constraint_forward(PyObject *self, PyObject *args)
-{
-    PyObject *prob_capsule;
-    PyObject *u_obj;
-    if (!PyArg_ParseTuple(args, "OO", &prob_capsule, &u_obj))
-    {
-        return NULL;
-    }
-
-    problem *prob =
-        (problem *) PyCapsule_GetPointer(prob_capsule, PROBLEM_CAPSULE_NAME);
-    if (!prob)
-    {
-        PyErr_SetString(PyExc_ValueError, "invalid problem capsule");
-        return NULL;
-    }
-
-    PyArrayObject *u_array =
-        (PyArrayObject *) PyArray_FROM_OTF(u_obj, NPY_DOUBLE, NPY_ARRAY_IN_ARRAY);
-    if (!u_array)
-    {
-        return NULL;
-    }
-
-    double *constraint_vals =
-        problem_constraint_forward(prob, (const double *) PyArray_DATA(u_array));
-
-    PyObject *out = NULL;
-    if (prob->total_constraint_size > 0)
-    {
-        npy_intp size = prob->total_constraint_size;
-        out = PyArray_SimpleNew(1, &size, NPY_DOUBLE);
-        if (!out)
-        {
-            Py_DECREF(u_array);
-            return NULL;
-        }
-        memcpy(PyArray_DATA((PyArrayObject *) out), constraint_vals,
-               size * sizeof(double));
-    }
-    else
-    {
-        npy_intp size = 0;
-        out = PyArray_SimpleNew(1, &size, NPY_DOUBLE);
-    }
-
-    Py_DECREF(u_array);
-    return out;
-}
-
-static PyObject *py_problem_gradient(PyObject *self, PyObject *args)
-{
-    PyObject *prob_capsule;
-    PyObject *u_obj;
-    if (!PyArg_ParseTuple(args, "OO", &prob_capsule, &u_obj))
-    {
-        return NULL;
-    }
-
-    problem *prob =
-        (problem *) PyCapsule_GetPointer(prob_capsule, PROBLEM_CAPSULE_NAME);
-    if (!prob)
-    {
-        PyErr_SetString(PyExc_ValueError, "invalid problem capsule");
-        return NULL;
-    }
-
-    PyArrayObject *u_array =
-        (PyArrayObject *) PyArray_FROM_OTF(u_obj, NPY_DOUBLE, NPY_ARRAY_IN_ARRAY);
-    if (!u_array)
-    {
-        return NULL;
-    }
-
-    double *grad = problem_gradient(prob, (const double *) PyArray_DATA(u_array));
-
-    npy_intp size = prob->n_vars;
-    PyObject *out = PyArray_SimpleNew(1, &size, NPY_DOUBLE);
-    if (!out)
-    {
-        Py_DECREF(u_array);
-        return NULL;
-    }
-    memcpy(PyArray_DATA((PyArrayObject *) out), grad, size * sizeof(double));
-
-    Py_DECREF(u_array);
-    return out;
-}
-
-static PyObject *py_problem_jacobian(PyObject *self, PyObject *args)
-{
-    PyObject *prob_capsule;
-    PyObject *u_obj;
-    if (!PyArg_ParseTuple(args, "OO", &prob_capsule, &u_obj))
-    {
-        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->n_constraints == 0)
-    {
-        // Return empty CSR components
-        npy_intp zero = 0;
-        npy_intp one = 1;
-        PyObject *data = PyArray_SimpleNew(1, &zero, NPY_DOUBLE);
-        PyObject *indices = PyArray_SimpleNew(1, &zero, NPY_INT32);
-        PyObject *indptr = PyArray_SimpleNew(1, &one, NPY_INT32);
-        ((int *) PyArray_DATA((PyArrayObject *) indptr))[0] = 0;
-        return Py_BuildValue("(OOO(ii))", data, indices, indptr, 0, prob->n_vars);
-    }
-
-    PyArrayObject *u_array =
-        (PyArrayObject *) PyArray_FROM_OTF(u_obj, NPY_DOUBLE, NPY_ARRAY_IN_ARRAY);
-    if (!u_array)
-    {
-        return NULL;
-    }
-
-    CSR_Matrix *jac = problem_jacobian(prob, (const double *) PyArray_DATA(u_array));
-
-    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);
-        Py_DECREF(u_array);
-        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));
-
-    Py_DECREF(u_array);
-    return Py_BuildValue("(OOO(ii))", data, indices, indptr, jac->m, jac->n);
-}
-
 static PyMethodDef DNLPMethods[] = {
     {"make_variable", py_make_variable, METH_VARARGS, "Create variable node"},
     {"make_constant", py_make_constant, METH_VARARGS, "Create constant node"},
@@ -649,9 +41,6 @@ static PyMethodDef DNLPMethods[] = {
     {"make_sum", py_make_sum, METH_VARARGS, "Create sum node"},
     {"make_neg", py_make_neg, METH_VARARGS, "Create neg node"},
     {"make_promote", py_make_promote, METH_VARARGS, "Create promote node"},
-    {"forward", py_forward, METH_VARARGS, "Run forward pass and return values"},
-    {"jacobian", py_jacobian, METH_VARARGS,
-     "Compute jacobian and return CSR components"},
     {"make_problem", py_make_problem, METH_VARARGS, "Create problem from objective and constraints"},
     {"problem_init_derivatives", py_problem_init_derivatives, METH_VARARGS, "Initialize derivative structures"},
     {"problem_objective_forward", py_problem_objective_forward, METH_VARARGS, "Evaluate objective only"},
diff --git a/python/problem/common.h b/python/problem/common.h
new file mode 100644
index 0000000..43e44c7
--- /dev/null
+++ b/python/problem/common.h
@@ -0,0 +1,24 @@
+#ifndef PROBLEM_COMMON_H
+#define PROBLEM_COMMON_H
+
+#define PY_SSIZE_T_CLEAN
+#include <Python.h>
+#include <numpy/arrayobject.h>
+
+#include "problem.h"
+
+/* Also need expr types for capsule handling */
+#include "../atoms/common.h"
+
+#define PROBLEM_CAPSULE_NAME "DNLP_PROBLEM"
+
+static void problem_capsule_destructor(PyObject *capsule)
+{
+    problem *prob = (problem *) PyCapsule_GetPointer(capsule, PROBLEM_CAPSULE_NAME);
+    if (prob)
+    {
+        free_problem(prob);
+    }
+}
+
+#endif /* PROBLEM_COMMON_H */
diff --git a/python/problem/constraint_forward.h b/python/problem/constraint_forward.h
new file mode 100644
index 0000000..9ddc920
--- /dev/null
+++ b/python/problem/constraint_forward.h
@@ -0,0 +1,56 @@
+#ifndef PROBLEM_CONSTRAINT_FORWARD_H
+#define PROBLEM_CONSTRAINT_FORWARD_H
+
+#include "common.h"
+
+static PyObject *py_problem_constraint_forward(PyObject *self, PyObject *args)
+{
+    PyObject *prob_capsule;
+    PyObject *u_obj;
+    if (!PyArg_ParseTuple(args, "OO", &prob_capsule, &u_obj))
+    {
+        return NULL;
+    }
+
+    problem *prob =
+        (problem *) PyCapsule_GetPointer(prob_capsule, PROBLEM_CAPSULE_NAME);
+    if (!prob)
+    {
+        PyErr_SetString(PyExc_ValueError, "invalid problem capsule");
+        return NULL;
+    }
+
+    PyArrayObject *u_array =
+        (PyArrayObject *) PyArray_FROM_OTF(u_obj, NPY_DOUBLE, NPY_ARRAY_IN_ARRAY);
+    if (!u_array)
+    {
+        return NULL;
+    }
+
+    double *constraint_vals =
+        problem_constraint_forward(prob, (const double *) PyArray_DATA(u_array));
+
+    PyObject *out = NULL;
+    if (prob->total_constraint_size > 0)
+    {
+        npy_intp size = prob->total_constraint_size;
+        out = PyArray_SimpleNew(1, &size, NPY_DOUBLE);
+        if (!out)
+        {
+            Py_DECREF(u_array);
+            return NULL;
+        }
+        memcpy(PyArray_DATA((PyArrayObject *) out), constraint_vals,
+               size * sizeof(double));
+    }
+    else
+    {
+        npy_intp size = 0;
+        out = PyArray_SimpleNew(1, &size, NPY_DOUBLE);
+    }
+
+    Py_DECREF(u_array);
+    return out;
+}
+
+#endif /* PROBLEM_CONSTRAINT_FORWARD_H */
diff --git a/python/problem/gradient.h b/python/problem/gradient.h
new file mode 100644
index 0000000..395b234
--- /dev/null
+++ b/python/problem/gradient.h
@@ -0,0 +1,45 @@
+#ifndef PROBLEM_GRADIENT_H
+#define PROBLEM_GRADIENT_H
+
+#include "common.h"
+
+static PyObject *py_problem_gradient(PyObject *self, PyObject *args)
+{
+    PyObject *prob_capsule;
+    PyObject *u_obj;
+    if (!PyArg_ParseTuple(args, "OO", &prob_capsule, &u_obj))
+    {
+        return NULL;
+    }
+
+    problem *prob =
+        (problem *) PyCapsule_GetPointer(prob_capsule, PROBLEM_CAPSULE_NAME);
+    if (!prob)
+    {
+        PyErr_SetString(PyExc_ValueError, "invalid problem capsule");
+        return NULL;
+    }
+
+    PyArrayObject *u_array =
+        (PyArrayObject *) PyArray_FROM_OTF(u_obj, NPY_DOUBLE, NPY_ARRAY_IN_ARRAY);
+    if (!u_array)
+    {
+        return NULL;
+    }
+
+    double *grad = problem_gradient(prob, (const double *) PyArray_DATA(u_array));
+
+    npy_intp size = prob->n_vars;
+    PyObject *out = PyArray_SimpleNew(1, &size, NPY_DOUBLE);
+    if (!out)
+    {
+        Py_DECREF(u_array);
+        return NULL;
+    }
+    memcpy(PyArray_DATA((PyArrayObject *) out), grad, size * sizeof(double));
+
+    Py_DECREF(u_array);
+    return out;
+}
+
+#endif /* PROBLEM_GRADIENT_H */
diff --git a/python/problem/init_derivatives.h b/python/problem/init_derivatives.h
new file mode 100644
index 0000000..cc6298d
--- /dev/null
+++ b/python/problem/init_derivatives.h
@@ -0,0 +1,27 @@
+#ifndef PROBLEM_INIT_DERIVATIVES_H
+#define PROBLEM_INIT_DERIVATIVES_H
+
+#include "common.h"
+
+static PyObject *py_problem_init_derivatives(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;
+    }
+
+    problem_init_derivatives(prob);
+
+    Py_RETURN_NONE;
+}
+
+#endif /* PROBLEM_INIT_DERIVATIVES_H */
diff --git a/python/problem/jacobian.h b/python/problem/jacobian.h
new file mode 100644
index 0000000..02acbc7
--- /dev/null
+++ b/python/problem/jacobian.h
@@ -0,0 +1,68 @@
+#ifndef PROBLEM_JACOBIAN_H
+#define PROBLEM_JACOBIAN_H
+
+#include "common.h"
+
+static PyObject *py_problem_jacobian(PyObject *self, PyObject *args)
+{
+    PyObject *prob_capsule;
+    PyObject *u_obj;
+    if (!PyArg_ParseTuple(args, "OO", &prob_capsule, &u_obj))
+    {
+        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->n_constraints == 0)
+    {
+        /* Return empty CSR components */
+        npy_intp zero = 0;
+        npy_intp one = 1;
+        PyObject *data = PyArray_SimpleNew(1, &zero, NPY_DOUBLE);
+        PyObject *indices = PyArray_SimpleNew(1, &zero, NPY_INT32);
+        PyObject *indptr = PyArray_SimpleNew(1, &one, NPY_INT32);
+        ((int *) PyArray_DATA((PyArrayObject *) indptr))[0] = 0;
+        return Py_BuildValue("(OOO(ii))", data, indices, indptr, 0, prob->n_vars);
+    }
+
+    PyArrayObject *u_array =
+        (PyArrayObject *) PyArray_FROM_OTF(u_obj, NPY_DOUBLE, NPY_ARRAY_IN_ARRAY);
+    if (!u_array)
+    {
+        return NULL;
+    }
+
+    CSR_Matrix *jac = problem_jacobian(prob, (const double *) PyArray_DATA(u_array));
+
+    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);
+        Py_DECREF(u_array);
+        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));
+
+    Py_DECREF(u_array);
+    return Py_BuildValue("(OOO(ii))", data, indices, indptr, jac->m, jac->n);
+}
+
+#endif /* PROBLEM_JACOBIAN_H */
diff --git a/python/problem/make_problem.h b/python/problem/make_problem.h
new file mode 100644
index 0000000..c9a9a97
--- /dev/null
+++ b/python/problem/make_problem.h
@@ -0,0 +1,64 @@
+#ifndef PROBLEM_MAKE_H
+#define PROBLEM_MAKE_H
+
+#include "common.h"
+
+static PyObject *py_make_problem(PyObject *self, PyObject *args)
+{
+    PyObject *obj_capsule;
+    PyObject *constraints_list;
+    if (!PyArg_ParseTuple(args, "OO", &obj_capsule, &constraints_list))
+    {
+        return NULL;
+    }
+
+    expr *objective = (expr *) PyCapsule_GetPointer(obj_capsule, EXPR_CAPSULE_NAME);
+    if (!objective)
+    {
+        PyErr_SetString(PyExc_ValueError, "invalid objective capsule");
+        return NULL;
+    }
+
+    if (!PyList_Check(constraints_list))
+    {
+        PyErr_SetString(PyExc_TypeError, "constraints must be a list");
+        return NULL;
+    }
+
+    Py_ssize_t n_constraints = PyList_Size(constraints_list);
+    expr **constraints = NULL;
+    if (n_constraints > 0)
+    {
+        constraints = (expr **) malloc(n_constraints * sizeof(expr *));
+        if (!constraints)
+        {
+            PyErr_NoMemory();
+            return NULL;
+        }
+        for (Py_ssize_t i = 0; i < n_constraints; i++)
+        {
+            PyObject *c_capsule = PyList_GetItem(constraints_list, i);
+            constraints[i] =
+                (expr *) PyCapsule_GetPointer(c_capsule, EXPR_CAPSULE_NAME);
+            if (!constraints[i])
+            {
+                free(constraints);
+                PyErr_SetString(PyExc_ValueError, "invalid constraint capsule");
+                return NULL;
+            }
+        }
+    }
+
+    problem *prob = new_problem(objective, constraints, (int) n_constraints);
+    free(constraints);
+
+    if (!prob)
+    {
+        PyErr_SetString(PyExc_RuntimeError, "failed to create problem");
+        return NULL;
+    }
+
+    return PyCapsule_New(prob, PROBLEM_CAPSULE_NAME, problem_capsule_destructor);
+}
+
+#endif /* PROBLEM_MAKE_H */
diff --git a/python/problem/objective_forward.h b/python/problem/objective_forward.h
new file mode 100644
index 0000000..0b3d79f
--- /dev/null
+++ b/python/problem/objective_forward.h
@@ -0,0 +1,36 @@
+#ifndef PROBLEM_OBJECTIVE_FORWARD_H
+#define PROBLEM_OBJECTIVE_FORWARD_H
+
+#include "common.h"
+
+static PyObject *py_problem_objective_forward(PyObject *self, PyObject *args)
+{
+    PyObject *prob_capsule;
+    PyObject *u_obj;
+    if (!PyArg_ParseTuple(args, "OO", &prob_capsule, &u_obj))
+    {
+        return NULL;
+    }
+
+    problem *prob =
+        (problem *) PyCapsule_GetPointer(prob_capsule, PROBLEM_CAPSULE_NAME);
+    if (!prob)
+    {
+        PyErr_SetString(PyExc_ValueError, "invalid problem capsule");
+        return NULL;
+    }
+
+    PyArrayObject *u_array =
+        (PyArrayObject *) PyArray_FROM_OTF(u_obj, NPY_DOUBLE, NPY_ARRAY_IN_ARRAY);
+    if (!u_array)
+    {
+        return NULL;
+    }
+
+    double obj_val = problem_objective_forward(prob, (const double *) PyArray_DATA(u_array));
+
+    Py_DECREF(u_array);
+    return Py_BuildValue("d", obj_val);
+}
+
+#endif /* PROBLEM_OBJECTIVE_FORWARD_H */

From aedaafdcbb2b7937789683fbeeb51cc4ba117585 Mon Sep 17 00:00:00 2001
From: William Zijie Zhang <william@gridmatic.com>
Date: Tue, 6 Jan 2026 13:20:14 -0500
Subject: [PATCH 13/27] Change gradient, jacobian, constraint_forward to return
 void
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Values are now accessed directly from the problem struct:
- prob->constraint_values for constraint forward results
- prob->gradient_values for objective gradient
- prob->stacked_jac for constraint jacobian

Also updated Python bindings and tests to use the new API.

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
---
 include/problem.h                   | 10 +++--
 python/convert.py                   | 12 +++---
 python/problem/constraint_forward.h |  8 ++--
 python/problem/gradient.h           | 16 ++------
 python/problem/jacobian.h           | 15 ++-----
 python/tests/test_constrained.py    | 26 ++++++------
 python/tests/test_problem_native.py | 12 +++---
 python/tests/test_unconstrained.py  | 62 ++++++++++++++---------------
 src/problem.c                       | 46 ++++++++++-----------
 tests/problem/test_problem.h        | 26 +++++++-----
 10 files changed, 106 insertions(+), 127 deletions(-)

diff --git a/include/problem.h b/include/problem.h
index d86bf9a..3159938 100644
--- a/include/problem.h
+++ b/include/problem.h
@@ -12,9 +12,11 @@ typedef struct problem
     int n_vars;
     int total_constraint_size;
 
-    /* Allocated by problem_init_derivatives */
+    /* Allocated by new_problem */
     double *constraint_values;
     double *gradient_values;
+
+    /* Allocated by problem_init_derivatives */
     CSR_Matrix *stacked_jac;
 } problem;
 
@@ -24,8 +26,8 @@ void problem_init_derivatives(problem *prob);
 void free_problem(problem *prob);
 
 double problem_objective_forward(problem *prob, const double *u);
-double *problem_constraint_forward(problem *prob, const double *u);
-double *problem_gradient(problem *prob, const double *u);
-CSR_Matrix *problem_jacobian(problem *prob, const double *u);
+void problem_constraint_forward(problem *prob, const double *u);
+void problem_gradient(problem *prob);
+void problem_jacobian(problem *prob);
 
 #endif
diff --git a/python/convert.py b/python/convert.py
index 4c344ed..3037812 100644
--- a/python/convert.py
+++ b/python/convert.py
@@ -151,11 +151,11 @@ def constraint_forward(self, u: np.ndarray) -> np.ndarray:
         """Evaluate constraints only. Returns constraint_values array."""
         return diffengine.problem_constraint_forward(self._capsule, u)
 
-    def gradient(self, u: np.ndarray) -> np.ndarray:
-        """Compute gradient of objective. Returns gradient array."""
-        return diffengine.problem_gradient(self._capsule, u)
+    def gradient(self) -> np.ndarray:
+        """Compute gradient of objective. Call objective_forward first. Returns gradient array."""
+        return diffengine.problem_gradient(self._capsule)
 
-    def jacobian(self, u: np.ndarray) -> sparse.csr_matrix:
-        """Compute jacobian of constraints. Returns scipy CSR matrix."""
-        data, indices, indptr, shape = diffengine.problem_jacobian(self._capsule, u)
+    def jacobian(self) -> sparse.csr_matrix:
+        """Compute jacobian of constraints. Call constraint_forward first. Returns scipy CSR matrix."""
+        data, indices, indptr, shape = diffengine.problem_jacobian(self._capsule)
         return sparse.csr_matrix((data, indices, indptr), shape=shape)
diff --git a/python/problem/constraint_forward.h b/python/problem/constraint_forward.h
index 9ddc920..9449db2 100644
--- a/python/problem/constraint_forward.h
+++ b/python/problem/constraint_forward.h
@@ -27,8 +27,8 @@ static PyObject *py_problem_constraint_forward(PyObject *self, PyObject *args)
         return NULL;
     }
 
-    double *constraint_vals =
-        problem_constraint_forward(prob, (const double *) PyArray_DATA(u_array));
+    problem_constraint_forward(prob, (const double *) PyArray_DATA(u_array));
+    Py_DECREF(u_array);
 
     PyObject *out = NULL;
     if (prob->total_constraint_size > 0)
@@ -37,10 +37,9 @@ static PyObject *py_problem_constraint_forward(PyObject *self, PyObject *args)
         out = PyArray_SimpleNew(1, &size, NPY_DOUBLE);
         if (!out)
         {
-            Py_DECREF(u_array);
             return NULL;
         }
-        memcpy(PyArray_DATA((PyArrayObject *) out), constraint_vals,
+        memcpy(PyArray_DATA((PyArrayObject *) out), prob->constraint_values,
                size * sizeof(double));
     }
     else
@@ -49,7 +48,6 @@ static PyObject *py_problem_constraint_forward(PyObject *self, PyObject *args)
         out = PyArray_SimpleNew(1, &size, NPY_DOUBLE);
     }
 
-    Py_DECREF(u_array);
     return out;
 }
 
diff --git a/python/problem/gradient.h b/python/problem/gradient.h
index 395b234..20cb5dd 100644
--- a/python/problem/gradient.h
+++ b/python/problem/gradient.h
@@ -6,8 +6,7 @@
 static PyObject *py_problem_gradient(PyObject *self, PyObject *args)
 {
     PyObject *prob_capsule;
-    PyObject *u_obj;
-    if (!PyArg_ParseTuple(args, "OO", &prob_capsule, &u_obj))
+    if (!PyArg_ParseTuple(args, "O", &prob_capsule))
     {
         return NULL;
     }
@@ -20,25 +19,16 @@ static PyObject *py_problem_gradient(PyObject *self, PyObject *args)
         return NULL;
     }
 
-    PyArrayObject *u_array =
-        (PyArrayObject *) PyArray_FROM_OTF(u_obj, NPY_DOUBLE, NPY_ARRAY_IN_ARRAY);
-    if (!u_array)
-    {
-        return NULL;
-    }
-
-    double *grad = problem_gradient(prob, (const double *) PyArray_DATA(u_array));
+    problem_gradient(prob);
 
     npy_intp size = prob->n_vars;
     PyObject *out = PyArray_SimpleNew(1, &size, NPY_DOUBLE);
     if (!out)
     {
-        Py_DECREF(u_array);
         return NULL;
     }
-    memcpy(PyArray_DATA((PyArrayObject *) out), grad, size * sizeof(double));
+    memcpy(PyArray_DATA((PyArrayObject *) out), prob->gradient_values, size * sizeof(double));
 
-    Py_DECREF(u_array);
     return out;
 }
 
diff --git a/python/problem/jacobian.h b/python/problem/jacobian.h
index 02acbc7..69f240d 100644
--- a/python/problem/jacobian.h
+++ b/python/problem/jacobian.h
@@ -6,8 +6,7 @@
 static PyObject *py_problem_jacobian(PyObject *self, PyObject *args)
 {
     PyObject *prob_capsule;
-    PyObject *u_obj;
-    if (!PyArg_ParseTuple(args, "OO", &prob_capsule, &u_obj))
+    if (!PyArg_ParseTuple(args, "O", &prob_capsule))
     {
         return NULL;
     }
@@ -32,15 +31,9 @@ static PyObject *py_problem_jacobian(PyObject *self, PyObject *args)
         return Py_BuildValue("(OOO(ii))", data, indices, indptr, 0, prob->n_vars);
     }
 
-    PyArrayObject *u_array =
-        (PyArrayObject *) PyArray_FROM_OTF(u_obj, NPY_DOUBLE, NPY_ARRAY_IN_ARRAY);
-    if (!u_array)
-    {
-        return NULL;
-    }
-
-    CSR_Matrix *jac = problem_jacobian(prob, (const double *) PyArray_DATA(u_array));
+    problem_jacobian(prob);
 
+    CSR_Matrix *jac = prob->stacked_jac;
     npy_intp nnz = jac->nnz;
     npy_intp m_plus_1 = jac->m + 1;
 
@@ -53,7 +46,6 @@ static PyObject *py_problem_jacobian(PyObject *self, PyObject *args)
         Py_XDECREF(data);
         Py_XDECREF(indices);
         Py_XDECREF(indptr);
-        Py_DECREF(u_array);
         return NULL;
     }
 
@@ -61,7 +53,6 @@ static PyObject *py_problem_jacobian(PyObject *self, PyObject *args)
     memcpy(PyArray_DATA((PyArrayObject *) indices), jac->i, nnz * sizeof(int));
     memcpy(PyArray_DATA((PyArrayObject *) indptr), jac->p, m_plus_1 * sizeof(int));
 
-    Py_DECREF(u_array);
     return Py_BuildValue("(OOO(ii))", data, indices, indptr, jac->m, jac->n);
 }
 
diff --git a/python/tests/test_constrained.py b/python/tests/test_constrained.py
index 23962f0..9574d3d 100644
--- a/python/tests/test_constrained.py
+++ b/python/tests/test_constrained.py
@@ -17,14 +17,14 @@ def test_single_constraint():
     prob = C_problem(problem)
 
     u = np.array([1.0, 2.0, 3.0])
-    prob.allocate(u)
+    prob.init_derivatives()
 
     # Test constraint_forward: constr.expr = -log(x)
     constraint_vals = prob.constraint_forward(u)
     assert np.allclose(constraint_vals, -np.log(u))
 
     # Test jacobian: d/dx(-log(x)) = -1/x
-    jac = prob.jacobian(u)
+    jac = prob.jacobian()
     expected_jac = np.diag(-1.0 / u)
     assert np.allclose(jac.toarray(), expected_jac)
 
@@ -42,7 +42,7 @@ def test_two_constraints():
     prob = C_problem(problem)
 
     u = np.array([1.0, 2.0])
-    prob.allocate(u)
+    prob.init_derivatives()
 
     # Test constraint_forward: [-log(u), -exp(u)]
     expected_constraint_vals = np.concatenate([-np.log(u), -np.exp(u)])
@@ -50,7 +50,7 @@ def test_two_constraints():
     assert np.allclose(constraint_vals, expected_constraint_vals)
 
     # Test jacobian - stacked vertically
-    jac = prob.jacobian(u)
+    jac = prob.jacobian()
     assert jac.shape == (4, 2)
     expected_jac = np.vstack([np.diag(-1.0 / u), np.diag(-np.exp(u))])
     assert np.allclose(jac.toarray(), expected_jac)
@@ -70,7 +70,7 @@ def test_three_constraints_different_sizes():
     prob = C_problem(problem)
 
     u = np.array([1.0, 2.0, 3.0])
-    prob.allocate(u)
+    prob.init_derivatives()
 
     # Test constraint_forward
     expected_constraint_vals = np.concatenate([
@@ -82,7 +82,7 @@ def test_three_constraints_different_sizes():
     assert np.allclose(constraint_vals, expected_constraint_vals)
 
     # Test jacobian shape and values
-    jac = prob.jacobian(u)
+    jac = prob.jacobian()
     assert jac.shape == (7, 3)
     # First 3 rows: -diag(1/u), next 3 rows: -diag(exp(u)), last row: -1/u
     expected_jac = np.zeros((7, 3))
@@ -108,7 +108,7 @@ def test_multiple_variables():
     x_vals = np.array([1.0, 2.0])
     y_vals = np.array([0.5, 1.0])
     u = np.concatenate([x_vals, y_vals])
-    prob.allocate(u)
+    prob.init_derivatives()
 
     # Test constraint_forward
     expected_constraint_vals = np.concatenate([-np.log(x_vals), -np.exp(y_vals)])
@@ -116,7 +116,7 @@ def test_multiple_variables():
     assert np.allclose(constraint_vals, expected_constraint_vals)
 
     # Test jacobian
-    jac = prob.jacobian(u)
+    jac = prob.jacobian()
     assert jac.shape == (4, 4)
     expected_jac = np.zeros((4, 4))
     expected_jac[0, 0] = -1.0 / x_vals[0]
@@ -142,7 +142,7 @@ def test_larger_scale():
     prob = C_problem(problem)
 
     u = np.linspace(1.0, 5.0, n)
-    prob.allocate(u)
+    prob.init_derivatives()
 
     # Test constraint_forward
     expected_constraint_vals = np.concatenate([
@@ -155,7 +155,7 @@ def test_larger_scale():
     assert np.allclose(constraint_vals, expected_constraint_vals)
 
     # Test jacobian shape
-    jac = prob.jacobian(u)
+    jac = prob.jacobian()
     assert jac.shape == (n + n + 1 + 1, n)
 
 
@@ -169,16 +169,16 @@ def test_repeated_evaluations():
     prob = C_problem(problem)
 
     u1 = np.array([1.0, 2.0, 3.0])
-    prob.allocate(u1)
+    prob.init_derivatives()
 
     # First evaluation
     constraint_vals1 = prob.constraint_forward(u1)
-    jac1 = prob.jacobian(u1)
+    jac1 = prob.jacobian()
 
     # Second evaluation at different point
     u2 = np.array([2.0, 3.0, 4.0])
     constraint_vals2 = prob.constraint_forward(u2)
-    jac2 = prob.jacobian(u2)
+    jac2 = prob.jacobian()
 
     assert np.allclose(constraint_vals1, -np.exp(u1))
     assert np.allclose(constraint_vals2, -np.exp(u2))
diff --git a/python/tests/test_problem_native.py b/python/tests/test_problem_native.py
index 56d9842..2e3a310 100644
--- a/python/tests/test_problem_native.py
+++ b/python/tests/test_problem_native.py
@@ -58,7 +58,7 @@ def test_problem_gradient():
     diffengine.problem_init_derivatives(prob)
 
     diffengine.problem_objective_forward(prob, u)
-    grad = diffengine.problem_gradient(prob, u)
+    grad = diffengine.problem_gradient(prob)
     expected_grad = 1.0 / u
     assert np.allclose(grad, expected_grad)
 
@@ -76,7 +76,7 @@ def test_problem_jacobian():
     diffengine.problem_init_derivatives(prob)
 
     diffengine.problem_constraint_forward(prob, u)
-    data, indices, indptr, shape = diffengine.problem_jacobian(prob, u)
+    data, indices, indptr, shape = diffengine.problem_jacobian(prob)
     jac = sparse.csr_matrix((data, indices, indptr), shape=shape)
 
     expected_jac = np.diag(1.0 / u)
@@ -100,11 +100,11 @@ def test_problem_no_constraints():
     assert len(constraint_vals) == 0
 
     diffengine.problem_objective_forward(prob, u)
-    grad = diffengine.problem_gradient(prob, u)
+    grad = diffengine.problem_gradient(prob)
     assert np.allclose(grad, 1.0 / u)
 
     diffengine.problem_constraint_forward(prob, u)
-    data, indices, indptr, shape = diffengine.problem_jacobian(prob, u)
+    data, indices, indptr, shape = diffengine.problem_jacobian(prob)
     jac = sparse.csr_matrix((data, indices, indptr), shape=shape)
     assert jac.shape == (0, 3)
 
@@ -138,13 +138,13 @@ def test_problem_multiple_constraints():
 
     # Test gradient
     diffengine.problem_objective_forward(prob, u)
-    grad = diffengine.problem_gradient(prob, u)
+    grad = diffengine.problem_gradient(prob)
     expected_grad = 1.0 / u
     assert np.allclose(grad, expected_grad)
 
     # Test Jacobian
     diffengine.problem_constraint_forward(prob, u)
-    data, indices, indptr, shape = diffengine.problem_jacobian(prob, u)
+    data, indices, indptr, shape = diffengine.problem_jacobian(prob)
     jac = sparse.csr_matrix((data, indices, indptr), shape=shape)
     
     # Expected Jacobian: 
diff --git a/python/tests/test_unconstrained.py b/python/tests/test_unconstrained.py
index 6cacd0b..99994a6 100644
--- a/python/tests/test_unconstrained.py
+++ b/python/tests/test_unconstrained.py
@@ -19,7 +19,7 @@ def test_sum_log():
     prob = C_problem(problem)
 
     u = np.array([1.0, 2.0, 3.0, 4.0])
-    prob.allocate(u)
+    prob.init_derivatives()
 
     # Objective
     obj_val = prob.objective_forward(u)
@@ -27,7 +27,7 @@ def test_sum_log():
     assert np.allclose(obj_val, expected)
 
     # Gradient: d/dx sum(log(x)) = 1/x
-    grad = prob.gradient(u)
+    grad = prob.gradient()
     expected_grad = 1.0 / u
     assert np.allclose(grad, expected_grad)
 
@@ -39,7 +39,7 @@ def test_sum_exp():
     prob = C_problem(problem)
 
     u = np.array([0.0, 1.0, 2.0])
-    prob.allocate(u)
+    prob.init_derivatives()
 
     # Objective
     obj_val = prob.objective_forward(u)
@@ -47,7 +47,7 @@ def test_sum_exp():
     assert np.allclose(obj_val, expected)
 
     # Gradient: d/dx sum(exp(x)) = exp(x)
-    grad = prob.gradient(u)
+    grad = prob.gradient()
     expected_grad = np.exp(u)
     assert np.allclose(grad, expected_grad)
 
@@ -59,7 +59,7 @@ def test_variable_reuse():
     prob = C_problem(problem)
 
     u = np.array([1.0, 2.0])
-    prob.allocate(u)
+    prob.init_derivatives()
 
     # Objective
     obj_val = prob.objective_forward(u)
@@ -67,7 +67,7 @@ def test_variable_reuse():
     assert np.allclose(obj_val, expected)
 
     # Gradient: d/dx_i = 1/x_i + exp(x_i)
-    grad = prob.gradient(u)
+    grad = prob.gradient()
     expected_grad = 1.0 / u + np.exp(u)
     assert np.allclose(grad, expected_grad)
 
@@ -79,7 +79,7 @@ def test_variable_used_multiple_times():
     prob = C_problem(problem)
 
     u = np.array([1.0, 2.0, 3.0])
-    prob.allocate(u)
+    prob.init_derivatives()
 
     # Objective
     obj_val = prob.objective_forward(u)
@@ -87,7 +87,7 @@ def test_variable_used_multiple_times():
     assert np.allclose(obj_val, expected)
 
     # Gradient: d/dx_i = 2/x_i + exp(x_i)
-    grad = prob.gradient(u)
+    grad = prob.gradient()
     expected_grad = 2.0 / u + np.exp(u)
     assert np.allclose(grad, expected_grad)
 
@@ -99,7 +99,7 @@ def test_larger_variable():
     prob = C_problem(problem)
 
     u = np.linspace(1.0, 10.0, 100)
-    prob.allocate(u)
+    prob.init_derivatives()
 
     # Objective
     obj_val = prob.objective_forward(u)
@@ -107,7 +107,7 @@ def test_larger_variable():
     assert np.allclose(obj_val, expected)
 
     # Gradient
-    grad = prob.gradient(u)
+    grad = prob.gradient()
     expected_grad = 1.0 / u
     assert np.allclose(grad, expected_grad)
 
@@ -119,7 +119,7 @@ def test_matrix_variable():
     prob = C_problem(problem)
 
     u = np.arange(1.0, 13.0)  # 12 elements
-    prob.allocate(u)
+    prob.init_derivatives()
 
     # Objective
     obj_val = prob.objective_forward(u)
@@ -127,7 +127,7 @@ def test_matrix_variable():
     assert np.allclose(obj_val, expected)
 
     # Gradient
-    grad = prob.gradient(u)
+    grad = prob.gradient()
     expected_grad = 1.0 / u
     assert np.allclose(grad, expected_grad)
 
@@ -142,7 +142,7 @@ def test_two_variables_separate_ops():
     x_vals = np.array([1.0, 2.0])
     y_vals = np.array([0.5, 1.0])
     u = np.concatenate([x_vals, y_vals])
-    prob.allocate(u)
+    prob.init_derivatives()
 
     # Objective
     obj_val = prob.objective_forward(u)
@@ -150,7 +150,7 @@ def test_two_variables_separate_ops():
     assert np.allclose(obj_val, expected)
 
     # Gradient
-    grad = prob.gradient(u)
+    grad = prob.gradient()
     expected_grad = np.concatenate([1.0 / x_vals, np.exp(y_vals)])
     assert np.allclose(grad, expected_grad)
 
@@ -165,7 +165,7 @@ def test_two_variables_same_sum():
     x_vals = np.array([1.0, 2.0])
     y_vals = np.array([3.0, 4.0])
     u = np.concatenate([x_vals, y_vals])
-    prob.allocate(u)
+    prob.init_derivatives()
 
     # Objective
     obj_val = prob.objective_forward(u)
@@ -173,7 +173,7 @@ def test_two_variables_same_sum():
     assert np.allclose(obj_val, expected)
 
     # Gradient
-    grad = prob.gradient(u)
+    grad = prob.gradient()
     expected_grad = np.concatenate([1.0 / x_vals, 1.0 / y_vals])
     assert np.allclose(grad, expected_grad)
 
@@ -192,7 +192,7 @@ def test_mixed_sizes():
     b_vals = np.array([1.0, 2.0, 3.0, 4.0, 5.0])
     c_vals = np.array([1.0, 2.0, 3.0])
     u = np.concatenate([a_vals, b_vals, c_vals])
-    prob.allocate(u)
+    prob.init_derivatives()
 
     # Objective
     obj_val = prob.objective_forward(u)
@@ -200,7 +200,7 @@ def test_mixed_sizes():
     assert np.allclose(obj_val, expected)
 
     # Gradient
-    grad = prob.gradient(u)
+    grad = prob.gradient()
     expected_grad = 1.0 / u
     assert np.allclose(grad, expected_grad)
 
@@ -220,7 +220,7 @@ def test_multiple_variables_log_exp():
     c_vals = np.array([0.1, 0.2])
     d_vals = np.array([0.1, 0.1])
     u = np.concatenate([a_vals, b_vals, c_vals, d_vals])
-    prob.allocate(u)
+    prob.init_derivatives()
 
     # Objective
     obj_val = prob.objective_forward(u)
@@ -229,7 +229,7 @@ def test_multiple_variables_log_exp():
     assert np.allclose(obj_val, expected)
 
     # Gradient
-    grad = prob.gradient(u)
+    grad = prob.gradient()
     expected_grad = np.concatenate([
         1.0 / a_vals,
         1.0 / b_vals,
@@ -252,7 +252,7 @@ def test_three_variables_mixed():
     y_vals = np.array([0.5, 1.0])
     z_vals = np.array([2.0, 3.0])
     u = np.concatenate([x_vals, y_vals, z_vals])
-    prob.allocate(u)
+    prob.init_derivatives()
 
     # Objective
     obj_val = prob.objective_forward(u)
@@ -260,7 +260,7 @@ def test_three_variables_mixed():
     assert np.allclose(obj_val, expected)
 
     # Gradient
-    grad = prob.gradient(u)
+    grad = prob.gradient()
     expected_grad = np.concatenate([1.0 / x_vals, np.exp(y_vals), 1.0 / z_vals])
     assert np.allclose(grad, expected_grad)
 
@@ -272,16 +272,16 @@ def test_repeated_evaluations():
     prob = C_problem(problem)
 
     u1 = np.array([1.0, 2.0, 3.0])
-    prob.allocate(u1)
+    prob.init_derivatives()
 
     # First evaluation
     obj1 = prob.objective_forward(u1)
-    grad1 = prob.gradient(u1)
+    grad1 = prob.gradient()
 
     # Second evaluation
     u2 = np.array([2.0, 3.0, 4.0])
     obj2 = prob.objective_forward(u2)
-    grad2 = prob.gradient(u2)
+    grad2 = prob.gradient()
 
     assert np.allclose(obj1, np.sum(np.log(u1)))
     assert np.allclose(obj2, np.sum(np.log(u2)))
@@ -320,7 +320,7 @@ def test_repeated_evaluations():
 #     x_vals = np.array([1.0, 2.0])
 #     y_vals = np.array([0.5, 1.0])
 #     u = np.concatenate([x_vals, y_vals])
-#     prob.allocate(u)
+#     prob.init_derivatives()
 #
 #     # Expected objective: sum(log(x + y))
 #     obj_val = prob.objective_forward(u)
@@ -328,7 +328,7 @@ def test_repeated_evaluations():
 #     assert np.allclose(obj_val, expected)
 #
 #     # Expected gradient: d/dx_i = 1/(x_i + y_i), d/dy_i = 1/(x_i + y_i)
-#     grad = prob.gradient(u)
+#     grad = prob.gradient()
 #     grad_xy = 1.0 / (x_vals + y_vals)
 #     expected_grad = np.concatenate([grad_xy, grad_xy])
 #     assert np.allclose(grad, expected_grad)
@@ -345,7 +345,7 @@ def test_repeated_evaluations():
 #     prob = C_problem(problem)
 #
 #     u = np.array([1.0, 2.0, 3.0, 4.0])
-#     prob.allocate(u)
+#     prob.init_derivatives()
 #
 #     # Expected objective: log(sum(x))
 #     obj_val = prob.objective_forward(u)
@@ -353,7 +353,7 @@ def test_repeated_evaluations():
 #     assert np.allclose(obj_val, expected)
 #
 #     # Expected gradient: d/dx_i log(sum(x)) = 1/sum(x)
-#     grad = prob.gradient(u)
+#     grad = prob.gradient()
 #     expected_grad = np.ones_like(u) / np.sum(u)
 #     assert np.allclose(grad, expected_grad)
 
@@ -374,7 +374,7 @@ def test_repeated_evaluations():
 #     y_vals = np.array([0.5, 1.0])
 #     z_vals = np.array([0.1, 0.2])
 #     u = np.concatenate([x_vals, y_vals, z_vals])
-#     prob.allocate(u)
+#     prob.init_derivatives()
 #
 #     # Expected objective
 #     obj_val = prob.objective_forward(u)
@@ -397,7 +397,7 @@ def test_repeated_evaluations():
 #     y_vals = np.array([0.5, 1.0])
 #     z_vals = np.array([0.5, 0.5])
 #     u = np.concatenate([x_vals, y_vals, z_vals])
-#     prob.allocate(u)
+#     prob.init_derivatives()
 #
 #     # Expected objective
 #     obj_val = prob.objective_forward(u)
diff --git a/src/problem.c b/src/problem.c
index 63281b5..b8c9e21 100644
--- a/src/problem.c
+++ b/src/problem.c
@@ -36,9 +36,18 @@ problem *new_problem(expr *objective, expr **constraints, int n_constraints)
 
     prob->n_vars = objective->n_vars;
 
-    /* Initialize allocated pointers to NULL */
-    prob->constraint_values = NULL;
-    prob->gradient_values = NULL;
+    /* Allocate value arrays */
+    if (prob->total_constraint_size > 0)
+    {
+        prob->constraint_values = (double *) calloc(prob->total_constraint_size, sizeof(double));
+    }
+    else
+    {
+        prob->constraint_values = NULL;
+    }
+    prob->gradient_values = (double *) calloc(prob->n_vars, sizeof(double));
+
+    /* Derivative structures allocated by problem_init_derivatives */
     prob->stacked_jac = NULL;
 
     return prob;
@@ -46,19 +55,10 @@ problem *new_problem(expr *objective, expr **constraints, int n_constraints)
 
 void problem_init_derivatives(problem *prob)
 {
-    /* 1. Allocate constraint values array */
-    if (prob->total_constraint_size > 0)
-    {
-        prob->constraint_values = (double *) calloc(prob->total_constraint_size, sizeof(double));
-    }
-
-    /* 2. Allocate gradient values array */
-    prob->gradient_values = (double *) calloc(prob->n_vars, sizeof(double));
-
-    /* 3. Initialize objective jacobian */
+    /* 1. Initialize objective jacobian */
     prob->objective->jacobian_init(prob->objective);
 
-    /* 4. Initialize constraint jacobians and count total nnz */
+    /* 2. Initialize constraint jacobians and count total nnz */
     int total_nnz = 0;
     for (int i = 0; i < prob->n_constraints; i++)
     {
@@ -67,15 +67,15 @@ void problem_init_derivatives(problem *prob)
         total_nnz += c->jacobian->nnz;
     }
 
-    /* 5. Allocate stacked jacobian */
+    /* 3. Allocate stacked jacobian */
     if (prob->total_constraint_size > 0)
     {
         prob->stacked_jac = new_csr_matrix(prob->total_constraint_size, prob->n_vars, total_nnz);
     }
 
-    /* TODO: 6. Initialize objective wsum_hess */
+    /* TODO: 4. Initialize objective wsum_hess */
 
-    /* TODO: 7. Initialize constraint wsum_hess */
+    /* TODO: 5. Initialize constraint wsum_hess */
 }
 
 void free_problem(problem *prob)
@@ -106,7 +106,7 @@ double problem_objective_forward(problem *prob, const double *u)
     return prob->objective->value[0];
 }
 
-double *problem_constraint_forward(problem *prob, const double *u)
+void problem_constraint_forward(problem *prob, const double *u)
 {
     /* Evaluate constraints only and copy values */
     int offset = 0;
@@ -117,11 +117,9 @@ double *problem_constraint_forward(problem *prob, const double *u)
         memcpy(prob->constraint_values + offset, c->value, c->size * sizeof(double));
         offset += c->size;
     }
-
-    return prob->constraint_values;
 }
 
-double *problem_gradient(problem *prob, const double *u)
+void problem_gradient(problem *prob)
 {
     /* Jacobian on objective */
     prob->objective->eval_jacobian(prob->objective);
@@ -137,11 +135,9 @@ double *problem_gradient(problem *prob, const double *u)
         int col = jac->i[k];
         prob->gradient_values[col] = jac->x[k];
     }
-
-    return prob->gradient_values;
 }
 
-CSR_Matrix *problem_jacobian(problem *prob, const double *u)
+void problem_jacobian(problem *prob)
 {
     CSR_Matrix *stacked = prob->stacked_jac;
 
@@ -178,6 +174,4 @@ CSR_Matrix *problem_jacobian(problem *prob, const double *u)
 
     /* Update actual nnz (may be less than allocated) */
     stacked->nnz = nnz_offset;
-
-    return stacked;
 }
diff --git a/tests/problem/test_problem.h b/tests/problem/test_problem.h
index b5e1c90..d4cc67a 100644
--- a/tests/problem/test_problem.h
+++ b/tests/problem/test_problem.h
@@ -117,15 +117,15 @@ const char *test_problem_constraint_forward(void)
     double u[2] = {2.0, 4.0};
     problem_init_derivatives(prob);
 
-    double *constraint_vals = problem_constraint_forward(prob, u);
+    problem_constraint_forward(prob, u);
 
     /* Check constraint values:
      * [log(2), log(4), exp(2), exp(4)]
      */
-    mu_assert("constraint[0] wrong", fabs(constraint_vals[0] - log(2.0)) < 1e-10);
-    mu_assert("constraint[1] wrong", fabs(constraint_vals[1] - log(4.0)) < 1e-10);
-    mu_assert("constraint[2] wrong", fabs(constraint_vals[2] - exp(2.0)) < 1e-10);
-    mu_assert("constraint[3] wrong", fabs(constraint_vals[3] - exp(4.0)) < 1e-10);
+    mu_assert("constraint[0] wrong", fabs(prob->constraint_values[0] - log(2.0)) < 1e-10);
+    mu_assert("constraint[1] wrong", fabs(prob->constraint_values[1] - log(4.0)) < 1e-10);
+    mu_assert("constraint[2] wrong", fabs(prob->constraint_values[2] - exp(2.0)) < 1e-10);
+    mu_assert("constraint[3] wrong", fabs(prob->constraint_values[3] - exp(4.0)) < 1e-10);
 
     free_problem(prob);
     free_expr(objective);
@@ -152,12 +152,12 @@ const char *test_problem_gradient(void)
     problem_init_derivatives(prob);
 
     problem_objective_forward(prob, u);
-    double *grad = problem_gradient(prob, u);
+    problem_gradient(prob);
 
     /* Expected gradient: [1/1, 1/2, 1/4] = [1.0, 0.5, 0.25] */
-    mu_assert("grad[0] wrong", fabs(grad[0] - 1.0) < 1e-10);
-    mu_assert("grad[1] wrong", fabs(grad[1] - 0.5) < 1e-10);
-    mu_assert("grad[2] wrong", fabs(grad[2] - 0.25) < 1e-10);
+    mu_assert("grad[0] wrong", fabs(prob->gradient_values[0] - 1.0) < 1e-10);
+    mu_assert("grad[1] wrong", fabs(prob->gradient_values[1] - 0.5) < 1e-10);
+    mu_assert("grad[2] wrong", fabs(prob->gradient_values[2] - 0.25) < 1e-10);
 
     free_problem(prob);
     free_expr(objective);
@@ -191,7 +191,9 @@ const char *test_problem_jacobian(void)
     problem_init_derivatives(prob);
 
     problem_constraint_forward(prob, u);
-    CSR_Matrix *jac = problem_jacobian(prob, u);
+    problem_jacobian(prob);
+
+    CSR_Matrix *jac = prob->stacked_jac;
 
     /* Check dimensions */
     mu_assert("jac rows wrong", jac->m == 2);
@@ -259,7 +261,9 @@ const char *test_problem_jacobian_multi(void)
     problem_init_derivatives(prob);
 
     problem_constraint_forward(prob, u);
-    CSR_Matrix *jac = problem_jacobian(prob, u);
+    problem_jacobian(prob);
+
+    CSR_Matrix *jac = prob->stacked_jac;
 
     /* Check dimensions: 4 rows (2 + 2), 2 cols */
     mu_assert("jac rows wrong", jac->m == 4);

From dad7950978568f0ff4d5b79987af69af31e0bb9d Mon Sep 17 00:00:00 2001
From: William Zijie Zhang <william@gridmatic.com>
Date: Tue, 6 Jan 2026 13:23:33 -0500
Subject: [PATCH 14/27] run formatting

---
 src/problem.c                |  6 ++++--
 tests/all_tests.c            |  2 +-
 tests/problem/test_problem.h | 12 ++++++++----
 3 files changed, 13 insertions(+), 7 deletions(-)

diff --git a/src/problem.c b/src/problem.c
index b8c9e21..6d9dec8 100644
--- a/src/problem.c
+++ b/src/problem.c
@@ -39,7 +39,8 @@ problem *new_problem(expr *objective, expr **constraints, int n_constraints)
     /* Allocate value arrays */
     if (prob->total_constraint_size > 0)
     {
-        prob->constraint_values = (double *) calloc(prob->total_constraint_size, sizeof(double));
+        prob->constraint_values =
+            (double *) calloc(prob->total_constraint_size, sizeof(double));
     }
     else
     {
@@ -70,7 +71,8 @@ void problem_init_derivatives(problem *prob)
     /* 3. Allocate stacked jacobian */
     if (prob->total_constraint_size > 0)
     {
-        prob->stacked_jac = new_csr_matrix(prob->total_constraint_size, prob->n_vars, total_nnz);
+        prob->stacked_jac =
+            new_csr_matrix(prob->total_constraint_size, prob->n_vars, total_nnz);
     }
 
     /* TODO: 4. Initialize objective wsum_hess */
diff --git a/tests/all_tests.c b/tests/all_tests.c
index fb3ed5f..6b26da9 100644
--- a/tests/all_tests.c
+++ b/tests/all_tests.c
@@ -21,6 +21,7 @@
 #include "jacobian_tests/test_quad_over_lin.h"
 #include "jacobian_tests/test_rel_entr.h"
 #include "jacobian_tests/test_sum.h"
+#include "problem/test_problem.h"
 #include "utils/test_csc_matrix.h"
 #include "utils/test_csr_matrix.h"
 #include "wsum_hess/elementwise/test_entr.h"
@@ -35,7 +36,6 @@
 #include "wsum_hess/test_multiply.h"
 #include "wsum_hess/test_rel_entr.h"
 #include "wsum_hess/test_sum.h"
-#include "problem/test_problem.h"
 
 int main(void)
 {
diff --git a/tests/problem/test_problem.h b/tests/problem/test_problem.h
index d4cc67a..1d42d00 100644
--- a/tests/problem/test_problem.h
+++ b/tests/problem/test_problem.h
@@ -122,10 +122,14 @@ const char *test_problem_constraint_forward(void)
     /* Check constraint values:
      * [log(2), log(4), exp(2), exp(4)]
      */
-    mu_assert("constraint[0] wrong", fabs(prob->constraint_values[0] - log(2.0)) < 1e-10);
-    mu_assert("constraint[1] wrong", fabs(prob->constraint_values[1] - log(4.0)) < 1e-10);
-    mu_assert("constraint[2] wrong", fabs(prob->constraint_values[2] - exp(2.0)) < 1e-10);
-    mu_assert("constraint[3] wrong", fabs(prob->constraint_values[3] - exp(4.0)) < 1e-10);
+    mu_assert("constraint[0] wrong",
+              fabs(prob->constraint_values[0] - log(2.0)) < 1e-10);
+    mu_assert("constraint[1] wrong",
+              fabs(prob->constraint_values[1] - log(4.0)) < 1e-10);
+    mu_assert("constraint[2] wrong",
+              fabs(prob->constraint_values[2] - exp(2.0)) < 1e-10);
+    mu_assert("constraint[3] wrong",
+              fabs(prob->constraint_values[3] - exp(4.0)) < 1e-10);
 
     free_problem(prob);
     free_expr(objective);

From 643fd2a5b0cd9c09d169aa96356f4c34696ebf4e Mon Sep 17 00:00:00 2001
From: William Zijie Zhang <william@gridmatic.com>
Date: Tue, 6 Jan 2026 13:29:28 -0500
Subject: [PATCH 15/27] run formatting and add venv to gitignore

---
 .gitignore                         |  5 ++++-
 python/.gitignore                  |  1 -
 python/atoms/constant.h            |  4 ++--
 python/atoms/linear.h              |  8 ++++----
 python/bindings.c                  | 32 ++++++++++++++++++------------
 python/problem/gradient.h          |  3 ++-
 python/problem/objective_forward.h |  3 ++-
 7 files changed, 33 insertions(+), 23 deletions(-)
 delete mode 100644 python/.gitignore

diff --git a/.gitignore b/.gitignore
index 7179491..f7de229 100644
--- a/.gitignore
+++ b/.gitignore
@@ -72,4 +72,7 @@ out/
 *.swp
 *.swo
 *~
-.DS_Store
\ No newline at end of file
+.DS_Store
+
+.venv/
+__pycache__/
\ No newline at end of file
diff --git a/python/.gitignore b/python/.gitignore
deleted file mode 100644
index ba0430d..0000000
--- a/python/.gitignore
+++ /dev/null
@@ -1 +0,0 @@
-__pycache__/
\ No newline at end of file
diff --git a/python/atoms/constant.h b/python/atoms/constant.h
index 8062647..da2af09 100644
--- a/python/atoms/constant.h
+++ b/python/atoms/constant.h
@@ -12,8 +12,8 @@ static PyObject *py_make_constant(PyObject *self, PyObject *args)
         return NULL;
     }
 
-    PyArrayObject *values_array =
-        (PyArrayObject *) PyArray_FROM_OTF(values_obj, NPY_DOUBLE, NPY_ARRAY_IN_ARRAY);
+    PyArrayObject *values_array = (PyArrayObject *) PyArray_FROM_OTF(
+        values_obj, NPY_DOUBLE, NPY_ARRAY_IN_ARRAY);
     if (!values_array)
     {
         return NULL;
diff --git a/python/atoms/linear.h b/python/atoms/linear.h
index 269b34f..0f85b91 100644
--- a/python/atoms/linear.h
+++ b/python/atoms/linear.h
@@ -23,10 +23,10 @@ static PyObject *py_make_linear(PyObject *self, PyObject *args)
 
     PyArrayObject *data_array =
         (PyArrayObject *) PyArray_FROM_OTF(data_obj, NPY_DOUBLE, NPY_ARRAY_IN_ARRAY);
-    PyArrayObject *indices_array =
-        (PyArrayObject *) PyArray_FROM_OTF(indices_obj, NPY_INT32, NPY_ARRAY_IN_ARRAY);
-    PyArrayObject *indptr_array =
-        (PyArrayObject *) PyArray_FROM_OTF(indptr_obj, NPY_INT32, NPY_ARRAY_IN_ARRAY);
+    PyArrayObject *indices_array = (PyArrayObject *) PyArray_FROM_OTF(
+        indices_obj, NPY_INT32, NPY_ARRAY_IN_ARRAY);
+    PyArrayObject *indptr_array = (PyArrayObject *) PyArray_FROM_OTF(
+        indptr_obj, NPY_INT32, NPY_ARRAY_IN_ARRAY);
 
     if (!data_array || !indices_array || !indptr_array)
     {
diff --git a/python/bindings.c b/python/bindings.c
index e153e0e..197390f 100644
--- a/python/bindings.c
+++ b/python/bindings.c
@@ -3,23 +3,23 @@
 #include <numpy/arrayobject.h>
 
 /* Include atom bindings */
-#include "atoms/variable.h"
+#include "atoms/add.h"
 #include "atoms/constant.h"
+#include "atoms/exp.h"
 #include "atoms/linear.h"
 #include "atoms/log.h"
-#include "atoms/exp.h"
-#include "atoms/add.h"
-#include "atoms/sum.h"
 #include "atoms/neg.h"
 #include "atoms/promote.h"
+#include "atoms/sum.h"
+#include "atoms/variable.h"
 
 /* Include problem bindings */
-#include "problem/make_problem.h"
-#include "problem/init_derivatives.h"
-#include "problem/objective_forward.h"
 #include "problem/constraint_forward.h"
 #include "problem/gradient.h"
+#include "problem/init_derivatives.h"
 #include "problem/jacobian.h"
+#include "problem/make_problem.h"
+#include "problem/objective_forward.h"
 
 static int numpy_initialized = 0;
 
@@ -41,12 +41,18 @@ static PyMethodDef DNLPMethods[] = {
     {"make_sum", py_make_sum, METH_VARARGS, "Create sum node"},
     {"make_neg", py_make_neg, METH_VARARGS, "Create neg node"},
     {"make_promote", py_make_promote, METH_VARARGS, "Create promote node"},
-    {"make_problem", py_make_problem, METH_VARARGS, "Create problem from objective and constraints"},
-    {"problem_init_derivatives", py_problem_init_derivatives, METH_VARARGS, "Initialize derivative structures"},
-    {"problem_objective_forward", py_problem_objective_forward, METH_VARARGS, "Evaluate objective only"},
-    {"problem_constraint_forward", py_problem_constraint_forward, METH_VARARGS, "Evaluate constraints only"},
-    {"problem_gradient", py_problem_gradient, METH_VARARGS, "Compute objective gradient"},
-    {"problem_jacobian", py_problem_jacobian, METH_VARARGS, "Compute constraint jacobian"},
+    {"make_problem", py_make_problem, METH_VARARGS,
+     "Create problem from objective and constraints"},
+    {"problem_init_derivatives", py_problem_init_derivatives, METH_VARARGS,
+     "Initialize derivative structures"},
+    {"problem_objective_forward", py_problem_objective_forward, METH_VARARGS,
+     "Evaluate objective only"},
+    {"problem_constraint_forward", py_problem_constraint_forward, METH_VARARGS,
+     "Evaluate constraints only"},
+    {"problem_gradient", py_problem_gradient, METH_VARARGS,
+     "Compute objective gradient"},
+    {"problem_jacobian", py_problem_jacobian, METH_VARARGS,
+     "Compute constraint jacobian"},
     {NULL, NULL, 0, NULL}};
 
 static struct PyModuleDef dnlp_module = {PyModuleDef_HEAD_INIT, "DNLP_diff_engine",
diff --git a/python/problem/gradient.h b/python/problem/gradient.h
index 20cb5dd..34b9b34 100644
--- a/python/problem/gradient.h
+++ b/python/problem/gradient.h
@@ -27,7 +27,8 @@ static PyObject *py_problem_gradient(PyObject *self, PyObject *args)
     {
         return NULL;
     }
-    memcpy(PyArray_DATA((PyArrayObject *) out), prob->gradient_values, size * sizeof(double));
+    memcpy(PyArray_DATA((PyArrayObject *) out), prob->gradient_values,
+           size * sizeof(double));
 
     return out;
 }
diff --git a/python/problem/objective_forward.h b/python/problem/objective_forward.h
index 0b3d79f..51b1060 100644
--- a/python/problem/objective_forward.h
+++ b/python/problem/objective_forward.h
@@ -27,7 +27,8 @@ static PyObject *py_problem_objective_forward(PyObject *self, PyObject *args)
         return NULL;
     }
 
-    double obj_val = problem_objective_forward(prob, (const double *) PyArray_DATA(u_array));
+    double obj_val =
+        problem_objective_forward(prob, (const double *) PyArray_DATA(u_array));
 
     Py_DECREF(u_array);
     return Py_BuildValue("d", obj_val);

From 25c5497854ae381f32ed66ccf23844a20f209572 Mon Sep 17 00:00:00 2001
From: William Zijie Zhang <william@gridmatic.com>
Date: Tue, 6 Jan 2026 14:45:50 -0500
Subject: [PATCH 16/27] Move sparsity pattern setup to init functions in neg
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

- Move row pointers and column indices copying from eval_jacobian to jacobian_init
- Move row pointers and column indices copying from eval_wsum_hess to wsum_hess_init
- Simplify eval_wsum_hess: pass weights directly to child and negate result
- Remove malloc/free and stdlib.h include

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
---
 src/affine/neg.c | 73 ++++++++++++++++++++----------------------------
 1 file changed, 30 insertions(+), 43 deletions(-)

diff --git a/src/affine/neg.c b/src/affine/neg.c
index 0af70ce..6ea0831 100644
--- a/src/affine/neg.c
+++ b/src/affine/neg.c
@@ -1,5 +1,4 @@
 #include "affine.h"
-#include <stdlib.h>
 
 static void forward(expr *node, const double *u)
 {
@@ -21,6 +20,17 @@ static void jacobian_init(expr *node)
     /* 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);
+
+    /* copy row pointers and column indices (sparsity pattern is constant) */
+    for (int i = 0; i <= child_jac->m; i++)
+    {
+        node->jacobian->p[i] = child_jac->p[i];
+    }
+    for (int k = 0; k < child_jac->nnz; k++)
+    {
+        node->jacobian->i[k] = child_jac->i[k];
+    }
+    node->jacobian->nnz = child_jac->nnz;
 }
 
 static void eval_jacobian(expr *node)
@@ -28,24 +38,12 @@ static void eval_jacobian(expr *node)
     /* evaluate child's jacobian */
     node->left->eval_jacobian(node->left);
 
-    /* negate child's jacobian: copy structure, negate values */
+    /* negate values only (sparsity pattern set in jacobian_init) */
     CSR_Matrix *child_jac = node->left->jacobian;
-    CSR_Matrix *jac = node->jacobian;
-
-    /* copy row pointers */
-    for (int i = 0; i <= child_jac->m; i++)
-    {
-        jac->p[i] = child_jac->p[i];
-    }
-
-    /* copy column indices and negate values */
-    int nnz = child_jac->p[child_jac->m];
-    for (int k = 0; k < nnz; k++)
+    for (int k = 0; k < child_jac->nnz; k++)
     {
-        jac->i[k] = child_jac->i[k];
-        jac->x[k] = -child_jac->x[k];
+        node->jacobian->x[k] = -child_jac->x[k];
     }
-    jac->nnz = nnz;
 }
 
 static void wsum_hess_init(expr *node)
@@ -56,41 +54,30 @@ static void wsum_hess_init(expr *node)
     /* 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);
-}
 
-static void eval_wsum_hess(expr *node, const double *w)
-{
-    /* For neg(x), d^2(-x)/dx^2 = 0, but we need to pass -w to child
-     * Actually: d/dx(-x) = -I, so Hessian contribution is (-I)^T H (-I) = H
-     * But the weight vector needs to be passed through: child sees same w */
-
-    /* Negate weights for child (chain rule for linear transformation) */
-    double *neg_w = (double *) malloc(node->size * sizeof(double));
-    for (int i = 0; i < node->size; i++)
+    /* copy row pointers and column indices (sparsity pattern is constant) */
+    for (int i = 0; i <= child_hess->m; i++)
     {
-        neg_w[i] = -w[i];
+        node->wsum_hess->p[i] = child_hess->p[i];
     }
-
-    /* evaluate child's wsum_hess with negated weights */
-    node->left->eval_wsum_hess(node->left, neg_w);
-    free(neg_w);
-
-    /* copy child's wsum_hess (the negation is already accounted for in weights) */
-    CSR_Matrix *child_hess = node->left->wsum_hess;
-    CSR_Matrix *hess = node->wsum_hess;
-
-    for (int i = 0; i <= child_hess->m; i++)
+    for (int k = 0; k < child_hess->nnz; k++)
     {
-        hess->p[i] = child_hess->p[i];
+        node->wsum_hess->i[k] = child_hess->i[k];
     }
+    node->wsum_hess->nnz = child_hess->nnz;
+}
+
+static void eval_wsum_hess(expr *node, const double *w)
+{
+    /* For f(x) = -g(x): d²f/dx² = -d²g/dx² */
+    node->left->eval_wsum_hess(node->left, w);
 
-    int nnz = child_hess->p[child_hess->m];
-    for (int k = 0; k < nnz; k++)
+    /* negate values (sparsity pattern set in wsum_hess_init) */
+    CSR_Matrix *child_hess = node->left->wsum_hess;
+    for (int k = 0; k < child_hess->nnz; k++)
     {
-        hess->i[k] = child_hess->i[k];
-        hess->x[k] = child_hess->x[k];
+        node->wsum_hess->x[k] = -child_hess->x[k];
     }
-    hess->nnz = nnz;
 }
 
 static bool is_affine(const expr *node)

From 2e7468a5e4807dd513ee831cc874313a661f0658 Mon Sep 17 00:00:00 2001
From: William Zijie Zhang <william@gridmatic.com>
Date: Tue, 6 Jan 2026 14:49:50 -0500
Subject: [PATCH 17/27] Move sparsity pattern setup to init functions in
 promote
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

- Move row pointers and column indices setup from eval_jacobian to jacobian_init
- Move row pointers and column indices copying from eval_wsum_hess to wsum_hess_init
- Simplify eval functions to only copy values

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
---
 src/affine/promote.c | 57 +++++++++++++++++++++++++-------------------
 1 file changed, 32 insertions(+), 25 deletions(-)

diff --git a/src/affine/promote.c b/src/affine/promote.c
index fcefc13..a40b7c9 100644
--- a/src/affine/promote.c
+++ b/src/affine/promote.c
@@ -37,6 +37,21 @@ static void jacobian_init(expr *node)
     if (nnz_max == 0) nnz_max = 1;
 
     node->jacobian = new_csr_matrix(node->size, node->n_vars, nnz_max);
+    CSR_Matrix *jac = node->jacobian;
+
+    /* Build sparsity pattern by replicating child's rows */
+    jac->nnz = 0;
+    for (int row = 0; row < node->size; row++)
+    {
+        jac->p[row] = jac->nnz;
+        int child_row = row % child_size;
+        for (int k = child_jac->p[child_row]; k < child_jac->p[child_row + 1]; k++)
+        {
+            jac->i[jac->nnz] = child_jac->i[k];
+            jac->nnz++;
+        }
+    }
+    jac->p[node->size] = jac->nnz;
 }
 
 static void eval_jacobian(expr *node)
@@ -48,24 +63,16 @@ static void eval_jacobian(expr *node)
     CSR_Matrix *jac = node->jacobian;
     int child_size = node->left->size;
 
-    /* Build jacobian by replicating child's rows */
-    jac->nnz = 0;
+    /* Copy values only (sparsity pattern set in jacobian_init) */
+    int idx = 0;
     for (int row = 0; row < node->size; row++)
     {
-        jac->p[row] = jac->nnz;
-
-        /* which child row does this output row correspond to? */
         int child_row = row % child_size;
-
-        /* copy entries from child_jac row */
         for (int k = child_jac->p[child_row]; k < child_jac->p[child_row + 1]; k++)
         {
-            jac->i[jac->nnz] = child_jac->i[k];
-            jac->x[jac->nnz] = child_jac->x[k];
-            jac->nnz++;
+            jac->x[idx++] = child_jac->x[k];
         }
     }
-    jac->p[node->size] = jac->nnz;
 }
 
 static void wsum_hess_init(expr *node)
@@ -76,6 +83,17 @@ static void wsum_hess_init(expr *node)
     /* same sparsity as child since we're summing weights */
     CSR_Matrix *child_hess = node->left->wsum_hess;
     node->wsum_hess = new_csr_matrix(child_hess->m, child_hess->n, child_hess->nnz);
+
+    /* copy row pointers and column indices (sparsity pattern is constant) */
+    for (int i = 0; i <= child_hess->m; i++)
+    {
+        node->wsum_hess->p[i] = child_hess->p[i];
+    }
+    for (int k = 0; k < child_hess->nnz; k++)
+    {
+        node->wsum_hess->i[k] = child_hess->i[k];
+    }
+    node->wsum_hess->nnz = child_hess->nnz;
 }
 
 static void eval_wsum_hess(expr *node, const double *w)
@@ -83,7 +101,6 @@ static void eval_wsum_hess(expr *node, const double *w)
     /* Sum weights that correspond to the same child element */
     int child_size = node->left->size;
     double *summed_w = (double *) calloc(child_size, sizeof(double));
-
     for (int i = 0; i < node->size; i++)
     {
         summed_w[i % child_size] += w[i];
@@ -93,22 +110,12 @@ static void eval_wsum_hess(expr *node, const double *w)
     node->left->eval_wsum_hess(node->left, summed_w);
     free(summed_w);
 
-    /* copy child's wsum_hess */
+    /* copy values only (sparsity pattern set in wsum_hess_init) */
     CSR_Matrix *child_hess = node->left->wsum_hess;
-    CSR_Matrix *hess = node->wsum_hess;
-
-    for (int i = 0; i <= child_hess->m; i++)
-    {
-        hess->p[i] = child_hess->p[i];
-    }
-
-    int nnz = child_hess->p[child_hess->m];
-    for (int k = 0; k < nnz; k++)
+    for (int k = 0; k < child_hess->nnz; k++)
     {
-        hess->i[k] = child_hess->i[k];
-        hess->x[k] = child_hess->x[k];
+        node->wsum_hess->x[k] = child_hess->x[k];
     }
-    hess->nnz = nnz;
 }
 
 static bool is_affine(const expr *node)

From 3701acf4817b6ce7109ced1fcd1cf531b373086b Mon Sep 17 00:00:00 2001
From: William Zijie Zhang <william@gridmatic.com>
Date: Tue, 6 Jan 2026 14:50:19 -0500
Subject: [PATCH 18/27] adds comment for convert

---
 python/convert.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/python/convert.py b/python/convert.py
index 3037812..92920cb 100644
--- a/python/convert.py
+++ b/python/convert.py
@@ -139,7 +139,7 @@ def __init__(self, cvxpy_problem: cp.Problem):
         self._allocated = False
 
     def init_derivatives(self):
-        """Initialize derivative structures. Must be called before gradient/jacobian."""
+        """Initialize derivative structures. Must be called before forward/gradient/jacobian."""
         diffengine.problem_init_derivatives(self._capsule)
         self._allocated = True
 

From 529c67c648cb720ed4fc9e9ea7559ce8c130048e Mon Sep 17 00:00:00 2001
From: William Zijie Zhang <william@gridmatic.com>
Date: Tue, 6 Jan 2026 15:06:00 -0500
Subject: [PATCH 19/27] Revert to debug CI failure

---
 src/affine/neg.c     | 73 ++++++++++++++++++++++++++------------------
 src/affine/promote.c | 57 +++++++++++++++-------------------
 2 files changed, 68 insertions(+), 62 deletions(-)

diff --git a/src/affine/neg.c b/src/affine/neg.c
index 6ea0831..0af70ce 100644
--- a/src/affine/neg.c
+++ b/src/affine/neg.c
@@ -1,4 +1,5 @@
 #include "affine.h"
+#include <stdlib.h>
 
 static void forward(expr *node, const double *u)
 {
@@ -20,17 +21,6 @@ static void jacobian_init(expr *node)
     /* 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);
-
-    /* copy row pointers and column indices (sparsity pattern is constant) */
-    for (int i = 0; i <= child_jac->m; i++)
-    {
-        node->jacobian->p[i] = child_jac->p[i];
-    }
-    for (int k = 0; k < child_jac->nnz; k++)
-    {
-        node->jacobian->i[k] = child_jac->i[k];
-    }
-    node->jacobian->nnz = child_jac->nnz;
 }
 
 static void eval_jacobian(expr *node)
@@ -38,12 +28,24 @@ static void eval_jacobian(expr *node)
     /* evaluate child's jacobian */
     node->left->eval_jacobian(node->left);
 
-    /* negate values only (sparsity pattern set in jacobian_init) */
+    /* negate child's jacobian: copy structure, negate values */
     CSR_Matrix *child_jac = node->left->jacobian;
-    for (int k = 0; k < child_jac->nnz; k++)
+    CSR_Matrix *jac = node->jacobian;
+
+    /* copy row pointers */
+    for (int i = 0; i <= child_jac->m; i++)
+    {
+        jac->p[i] = child_jac->p[i];
+    }
+
+    /* copy column indices and negate values */
+    int nnz = child_jac->p[child_jac->m];
+    for (int k = 0; k < nnz; k++)
     {
-        node->jacobian->x[k] = -child_jac->x[k];
+        jac->i[k] = child_jac->i[k];
+        jac->x[k] = -child_jac->x[k];
     }
+    jac->nnz = nnz;
 }
 
 static void wsum_hess_init(expr *node)
@@ -54,30 +56,41 @@ static void wsum_hess_init(expr *node)
     /* 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);
-
-    /* copy row pointers and column indices (sparsity pattern is constant) */
-    for (int i = 0; i <= child_hess->m; i++)
-    {
-        node->wsum_hess->p[i] = child_hess->p[i];
-    }
-    for (int k = 0; k < child_hess->nnz; k++)
-    {
-        node->wsum_hess->i[k] = child_hess->i[k];
-    }
-    node->wsum_hess->nnz = child_hess->nnz;
 }
 
 static void eval_wsum_hess(expr *node, const double *w)
 {
-    /* For f(x) = -g(x): d²f/dx² = -d²g/dx² */
-    node->left->eval_wsum_hess(node->left, w);
+    /* For neg(x), d^2(-x)/dx^2 = 0, but we need to pass -w to child
+     * Actually: d/dx(-x) = -I, so Hessian contribution is (-I)^T H (-I) = H
+     * But the weight vector needs to be passed through: child sees same w */
+
+    /* Negate weights for child (chain rule for linear transformation) */
+    double *neg_w = (double *) malloc(node->size * sizeof(double));
+    for (int i = 0; i < node->size; i++)
+    {
+        neg_w[i] = -w[i];
+    }
 
-    /* negate values (sparsity pattern set in wsum_hess_init) */
+    /* evaluate child's wsum_hess with negated weights */
+    node->left->eval_wsum_hess(node->left, neg_w);
+    free(neg_w);
+
+    /* copy child's wsum_hess (the negation is already accounted for in weights) */
     CSR_Matrix *child_hess = node->left->wsum_hess;
-    for (int k = 0; k < child_hess->nnz; k++)
+    CSR_Matrix *hess = node->wsum_hess;
+
+    for (int i = 0; i <= child_hess->m; i++)
+    {
+        hess->p[i] = child_hess->p[i];
+    }
+
+    int nnz = child_hess->p[child_hess->m];
+    for (int k = 0; k < nnz; k++)
     {
-        node->wsum_hess->x[k] = -child_hess->x[k];
+        hess->i[k] = child_hess->i[k];
+        hess->x[k] = child_hess->x[k];
     }
+    hess->nnz = nnz;
 }
 
 static bool is_affine(const expr *node)
diff --git a/src/affine/promote.c b/src/affine/promote.c
index a40b7c9..fcefc13 100644
--- a/src/affine/promote.c
+++ b/src/affine/promote.c
@@ -37,21 +37,6 @@ static void jacobian_init(expr *node)
     if (nnz_max == 0) nnz_max = 1;
 
     node->jacobian = new_csr_matrix(node->size, node->n_vars, nnz_max);
-    CSR_Matrix *jac = node->jacobian;
-
-    /* Build sparsity pattern by replicating child's rows */
-    jac->nnz = 0;
-    for (int row = 0; row < node->size; row++)
-    {
-        jac->p[row] = jac->nnz;
-        int child_row = row % child_size;
-        for (int k = child_jac->p[child_row]; k < child_jac->p[child_row + 1]; k++)
-        {
-            jac->i[jac->nnz] = child_jac->i[k];
-            jac->nnz++;
-        }
-    }
-    jac->p[node->size] = jac->nnz;
 }
 
 static void eval_jacobian(expr *node)
@@ -63,16 +48,24 @@ static void eval_jacobian(expr *node)
     CSR_Matrix *jac = node->jacobian;
     int child_size = node->left->size;
 
-    /* Copy values only (sparsity pattern set in jacobian_init) */
-    int idx = 0;
+    /* Build jacobian by replicating child's rows */
+    jac->nnz = 0;
     for (int row = 0; row < node->size; row++)
     {
+        jac->p[row] = jac->nnz;
+
+        /* which child row does this output row correspond to? */
         int child_row = row % child_size;
+
+        /* copy entries from child_jac row */
         for (int k = child_jac->p[child_row]; k < child_jac->p[child_row + 1]; k++)
         {
-            jac->x[idx++] = child_jac->x[k];
+            jac->i[jac->nnz] = child_jac->i[k];
+            jac->x[jac->nnz] = child_jac->x[k];
+            jac->nnz++;
         }
     }
+    jac->p[node->size] = jac->nnz;
 }
 
 static void wsum_hess_init(expr *node)
@@ -83,17 +76,6 @@ static void wsum_hess_init(expr *node)
     /* same sparsity as child since we're summing weights */
     CSR_Matrix *child_hess = node->left->wsum_hess;
     node->wsum_hess = new_csr_matrix(child_hess->m, child_hess->n, child_hess->nnz);
-
-    /* copy row pointers and column indices (sparsity pattern is constant) */
-    for (int i = 0; i <= child_hess->m; i++)
-    {
-        node->wsum_hess->p[i] = child_hess->p[i];
-    }
-    for (int k = 0; k < child_hess->nnz; k++)
-    {
-        node->wsum_hess->i[k] = child_hess->i[k];
-    }
-    node->wsum_hess->nnz = child_hess->nnz;
 }
 
 static void eval_wsum_hess(expr *node, const double *w)
@@ -101,6 +83,7 @@ static void eval_wsum_hess(expr *node, const double *w)
     /* Sum weights that correspond to the same child element */
     int child_size = node->left->size;
     double *summed_w = (double *) calloc(child_size, sizeof(double));
+
     for (int i = 0; i < node->size; i++)
     {
         summed_w[i % child_size] += w[i];
@@ -110,12 +93,22 @@ static void eval_wsum_hess(expr *node, const double *w)
     node->left->eval_wsum_hess(node->left, summed_w);
     free(summed_w);
 
-    /* copy values only (sparsity pattern set in wsum_hess_init) */
+    /* copy child's wsum_hess */
     CSR_Matrix *child_hess = node->left->wsum_hess;
-    for (int k = 0; k < child_hess->nnz; k++)
+    CSR_Matrix *hess = node->wsum_hess;
+
+    for (int i = 0; i <= child_hess->m; i++)
+    {
+        hess->p[i] = child_hess->p[i];
+    }
+
+    int nnz = child_hess->p[child_hess->m];
+    for (int k = 0; k < nnz; k++)
     {
-        node->wsum_hess->x[k] = child_hess->x[k];
+        hess->i[k] = child_hess->i[k];
+        hess->x[k] = child_hess->x[k];
     }
+    hess->nnz = nnz;
 }
 
 static bool is_affine(const expr *node)

From 1c3a1d574de72750c92b9e682531be21c1c59417 Mon Sep 17 00:00:00 2001
From: William Zijie Zhang <william@gridmatic.com>
Date: Tue, 6 Jan 2026 15:14:18 -0500
Subject: [PATCH 20/27] Move sparsity pattern to init functions and fix tests
 for new memory model
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

neg.c / promote.c:
- Move row pointers and column indices copying from eval to init functions
- Sparsity pattern is constant, only values change between iterations
- Simplify eval_wsum_hess in neg: pass weights directly and negate result

Tests:
- Update test_neg.h, test_promote.h, test_problem.h for refcount=0 model
- Only free root nodes; children are freed automatically via refcounting

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
---
 src/affine/neg.c                         | 73 ++++++++++--------------
 src/affine/promote.c                     | 57 ++++++++++--------
 tests/forward_pass/affine/test_neg.h     |  4 --
 tests/forward_pass/affine/test_promote.h |  3 -
 tests/problem/test_problem.h             | 28 +--------
 5 files changed, 63 insertions(+), 102 deletions(-)

diff --git a/src/affine/neg.c b/src/affine/neg.c
index 0af70ce..6ea0831 100644
--- a/src/affine/neg.c
+++ b/src/affine/neg.c
@@ -1,5 +1,4 @@
 #include "affine.h"
-#include <stdlib.h>
 
 static void forward(expr *node, const double *u)
 {
@@ -21,6 +20,17 @@ static void jacobian_init(expr *node)
     /* 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);
+
+    /* copy row pointers and column indices (sparsity pattern is constant) */
+    for (int i = 0; i <= child_jac->m; i++)
+    {
+        node->jacobian->p[i] = child_jac->p[i];
+    }
+    for (int k = 0; k < child_jac->nnz; k++)
+    {
+        node->jacobian->i[k] = child_jac->i[k];
+    }
+    node->jacobian->nnz = child_jac->nnz;
 }
 
 static void eval_jacobian(expr *node)
@@ -28,24 +38,12 @@ static void eval_jacobian(expr *node)
     /* evaluate child's jacobian */
     node->left->eval_jacobian(node->left);
 
-    /* negate child's jacobian: copy structure, negate values */
+    /* negate values only (sparsity pattern set in jacobian_init) */
     CSR_Matrix *child_jac = node->left->jacobian;
-    CSR_Matrix *jac = node->jacobian;
-
-    /* copy row pointers */
-    for (int i = 0; i <= child_jac->m; i++)
-    {
-        jac->p[i] = child_jac->p[i];
-    }
-
-    /* copy column indices and negate values */
-    int nnz = child_jac->p[child_jac->m];
-    for (int k = 0; k < nnz; k++)
+    for (int k = 0; k < child_jac->nnz; k++)
     {
-        jac->i[k] = child_jac->i[k];
-        jac->x[k] = -child_jac->x[k];
+        node->jacobian->x[k] = -child_jac->x[k];
     }
-    jac->nnz = nnz;
 }
 
 static void wsum_hess_init(expr *node)
@@ -56,41 +54,30 @@ static void wsum_hess_init(expr *node)
     /* 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);
-}
 
-static void eval_wsum_hess(expr *node, const double *w)
-{
-    /* For neg(x), d^2(-x)/dx^2 = 0, but we need to pass -w to child
-     * Actually: d/dx(-x) = -I, so Hessian contribution is (-I)^T H (-I) = H
-     * But the weight vector needs to be passed through: child sees same w */
-
-    /* Negate weights for child (chain rule for linear transformation) */
-    double *neg_w = (double *) malloc(node->size * sizeof(double));
-    for (int i = 0; i < node->size; i++)
+    /* copy row pointers and column indices (sparsity pattern is constant) */
+    for (int i = 0; i <= child_hess->m; i++)
     {
-        neg_w[i] = -w[i];
+        node->wsum_hess->p[i] = child_hess->p[i];
     }
-
-    /* evaluate child's wsum_hess with negated weights */
-    node->left->eval_wsum_hess(node->left, neg_w);
-    free(neg_w);
-
-    /* copy child's wsum_hess (the negation is already accounted for in weights) */
-    CSR_Matrix *child_hess = node->left->wsum_hess;
-    CSR_Matrix *hess = node->wsum_hess;
-
-    for (int i = 0; i <= child_hess->m; i++)
+    for (int k = 0; k < child_hess->nnz; k++)
     {
-        hess->p[i] = child_hess->p[i];
+        node->wsum_hess->i[k] = child_hess->i[k];
     }
+    node->wsum_hess->nnz = child_hess->nnz;
+}
+
+static void eval_wsum_hess(expr *node, const double *w)
+{
+    /* For f(x) = -g(x): d²f/dx² = -d²g/dx² */
+    node->left->eval_wsum_hess(node->left, w);
 
-    int nnz = child_hess->p[child_hess->m];
-    for (int k = 0; k < nnz; k++)
+    /* negate values (sparsity pattern set in wsum_hess_init) */
+    CSR_Matrix *child_hess = node->left->wsum_hess;
+    for (int k = 0; k < child_hess->nnz; k++)
     {
-        hess->i[k] = child_hess->i[k];
-        hess->x[k] = child_hess->x[k];
+        node->wsum_hess->x[k] = -child_hess->x[k];
     }
-    hess->nnz = nnz;
 }
 
 static bool is_affine(const expr *node)
diff --git a/src/affine/promote.c b/src/affine/promote.c
index fcefc13..a40b7c9 100644
--- a/src/affine/promote.c
+++ b/src/affine/promote.c
@@ -37,6 +37,21 @@ static void jacobian_init(expr *node)
     if (nnz_max == 0) nnz_max = 1;
 
     node->jacobian = new_csr_matrix(node->size, node->n_vars, nnz_max);
+    CSR_Matrix *jac = node->jacobian;
+
+    /* Build sparsity pattern by replicating child's rows */
+    jac->nnz = 0;
+    for (int row = 0; row < node->size; row++)
+    {
+        jac->p[row] = jac->nnz;
+        int child_row = row % child_size;
+        for (int k = child_jac->p[child_row]; k < child_jac->p[child_row + 1]; k++)
+        {
+            jac->i[jac->nnz] = child_jac->i[k];
+            jac->nnz++;
+        }
+    }
+    jac->p[node->size] = jac->nnz;
 }
 
 static void eval_jacobian(expr *node)
@@ -48,24 +63,16 @@ static void eval_jacobian(expr *node)
     CSR_Matrix *jac = node->jacobian;
     int child_size = node->left->size;
 
-    /* Build jacobian by replicating child's rows */
-    jac->nnz = 0;
+    /* Copy values only (sparsity pattern set in jacobian_init) */
+    int idx = 0;
     for (int row = 0; row < node->size; row++)
     {
-        jac->p[row] = jac->nnz;
-
-        /* which child row does this output row correspond to? */
         int child_row = row % child_size;
-
-        /* copy entries from child_jac row */
         for (int k = child_jac->p[child_row]; k < child_jac->p[child_row + 1]; k++)
         {
-            jac->i[jac->nnz] = child_jac->i[k];
-            jac->x[jac->nnz] = child_jac->x[k];
-            jac->nnz++;
+            jac->x[idx++] = child_jac->x[k];
         }
     }
-    jac->p[node->size] = jac->nnz;
 }
 
 static void wsum_hess_init(expr *node)
@@ -76,6 +83,17 @@ static void wsum_hess_init(expr *node)
     /* same sparsity as child since we're summing weights */
     CSR_Matrix *child_hess = node->left->wsum_hess;
     node->wsum_hess = new_csr_matrix(child_hess->m, child_hess->n, child_hess->nnz);
+
+    /* copy row pointers and column indices (sparsity pattern is constant) */
+    for (int i = 0; i <= child_hess->m; i++)
+    {
+        node->wsum_hess->p[i] = child_hess->p[i];
+    }
+    for (int k = 0; k < child_hess->nnz; k++)
+    {
+        node->wsum_hess->i[k] = child_hess->i[k];
+    }
+    node->wsum_hess->nnz = child_hess->nnz;
 }
 
 static void eval_wsum_hess(expr *node, const double *w)
@@ -83,7 +101,6 @@ static void eval_wsum_hess(expr *node, const double *w)
     /* Sum weights that correspond to the same child element */
     int child_size = node->left->size;
     double *summed_w = (double *) calloc(child_size, sizeof(double));
-
     for (int i = 0; i < node->size; i++)
     {
         summed_w[i % child_size] += w[i];
@@ -93,22 +110,12 @@ static void eval_wsum_hess(expr *node, const double *w)
     node->left->eval_wsum_hess(node->left, summed_w);
     free(summed_w);
 
-    /* copy child's wsum_hess */
+    /* copy values only (sparsity pattern set in wsum_hess_init) */
     CSR_Matrix *child_hess = node->left->wsum_hess;
-    CSR_Matrix *hess = node->wsum_hess;
-
-    for (int i = 0; i <= child_hess->m; i++)
-    {
-        hess->p[i] = child_hess->p[i];
-    }
-
-    int nnz = child_hess->p[child_hess->m];
-    for (int k = 0; k < nnz; k++)
+    for (int k = 0; k < child_hess->nnz; k++)
     {
-        hess->i[k] = child_hess->i[k];
-        hess->x[k] = child_hess->x[k];
+        node->wsum_hess->x[k] = child_hess->x[k];
     }
-    hess->nnz = nnz;
 }
 
 static bool is_affine(const expr *node)
diff --git a/tests/forward_pass/affine/test_neg.h b/tests/forward_pass/affine/test_neg.h
index 73c14e3..af0c9aa 100644
--- a/tests/forward_pass/affine/test_neg.h
+++ b/tests/forward_pass/affine/test_neg.h
@@ -16,7 +16,6 @@ const char *test_neg_forward(void)
     double expected[3] = {-1.0, -2.0, -3.0};
     mu_assert("neg forward failed", cmp_double_array(neg_node->value, expected, 3));
     free_expr(neg_node);
-    free_expr(var);
     return 0;
 }
 
@@ -42,7 +41,6 @@ const char *test_neg_jacobian(void)
               cmp_int_array(neg_node->jacobian->i, expected_i, 3));
 
     free_expr(neg_node);
-    free_expr(var);
     return 0;
 }
 
@@ -74,7 +72,5 @@ const char *test_neg_chain(void)
               cmp_int_array(neg2->jacobian->i, expected_i, 3));
 
     free_expr(neg2);
-    free_expr(neg1);
-    free_expr(var);
     return 0;
 }
diff --git a/tests/forward_pass/affine/test_promote.h b/tests/forward_pass/affine/test_promote.h
index eda5bf3..7195fb0 100644
--- a/tests/forward_pass/affine/test_promote.h
+++ b/tests/forward_pass/affine/test_promote.h
@@ -20,7 +20,6 @@ const char *test_promote_scalar_to_vector(void)
               cmp_double_array(promote_node->value, expected, 4));
 
     free_expr(promote_node);
-    free_expr(var);
     return 0;
 }
 
@@ -47,7 +46,6 @@ const char *test_promote_scalar_jacobian(void)
               cmp_int_array(promote_node->jacobian->i, expected_i, 3));
 
     free_expr(promote_node);
-    free_expr(var);
     return 0;
 }
 
@@ -85,6 +83,5 @@ const char *test_promote_vector_jacobian(void)
               cmp_int_array(promote_node->jacobian->i, expected_i, 4));
 
     free_expr(promote_node);
-    free_expr(var);
     return 0;
 }
diff --git a/tests/problem/test_problem.h b/tests/problem/test_problem.h
index 1d42d00..1c0e96b 100644
--- a/tests/problem/test_problem.h
+++ b/tests/problem/test_problem.h
@@ -37,12 +37,8 @@ const char *test_problem_new_free(void)
     mu_assert("n_constraints wrong", prob->n_constraints == 1);
     mu_assert("total_constraint_size wrong", prob->total_constraint_size == 3);
 
-    /* Free problem and expressions (shared ownership) */
+    /* Free problem (recursively frees expressions) */
     free_problem(prob);
-    free_expr(objective);
-    free_expr(log_x);
-    free_expr(x);
-    free_expr(x_constraint);
 
     return 0;
 }
@@ -80,10 +76,6 @@ const char *test_problem_objective_forward(void)
     mu_assert("constraint[2] wrong", fabs(prob->constraint_values[2] - 3.0) < 1e-10);
 
     free_problem(prob);
-    free_expr(objective);
-    free_expr(log_x);
-    free_expr(x);
-    free_expr(x_constraint);
 
     return 0;
 }
@@ -132,11 +124,6 @@ const char *test_problem_constraint_forward(void)
               fabs(prob->constraint_values[3] - exp(4.0)) < 1e-10);
 
     free_problem(prob);
-    free_expr(objective);
-    free_expr(log_obj);
-    free_expr(log_c1);
-    free_expr(exp_c2);
-    free_expr(x);
 
     return 0;
 }
@@ -164,9 +151,6 @@ const char *test_problem_gradient(void)
     mu_assert("grad[2] wrong", fabs(prob->gradient_values[2] - 0.25) < 1e-10);
 
     free_problem(prob);
-    free_expr(objective);
-    free_expr(log_x);
-    free_expr(x);
 
     return 0;
 }
@@ -217,11 +201,6 @@ const char *test_problem_jacobian(void)
     mu_assert("jac->x[1] wrong", fabs(jac->x[1] - 0.25) < 1e-10);
 
     free_problem(prob);
-    free_expr(objective);
-    free_expr(log_obj);
-    free_expr(x_obj);
-    free_expr(log_c1);
-    free_expr(x_c1);
 
     return 0;
 }
@@ -299,11 +278,6 @@ const char *test_problem_jacobian_multi(void)
     mu_assert("jac->x[3] wrong", fabs(jac->x[3] - exp(4.0)) < 1e-10);
 
     free_problem(prob);
-    free_expr(objective);
-    free_expr(log_obj);
-    free_expr(log_c1);
-    free_expr(exp_c2);
-    free_expr(x);
 
     return 0;
 }

From 26c0a1a70bbbb60a0ca3beec58fa252dbd185109 Mon Sep 17 00:00:00 2001
From: William Zijie Zhang <william@gridmatic.com>
Date: Wed, 7 Jan 2026 15:48:21 -0500
Subject: [PATCH 21/27] Use memcpy for CSR sparsity copies and move allocation
 to init
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

- Replace manual loops with memcpy for copying row pointers and column
  indices in neg.c and promote.c
- Move workspace allocation from eval_wsum_hess to wsum_hess_init in
  promote.c, using memset to zero before use
- Remove unused promote_expr struct (had no extra fields beyond base)
- Simplify new_promote to use new_expr directly

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
---
 include/subexpr.h    |  7 -------
 src/affine/neg.c     | 21 +++++----------------
 src/affine/promote.c | 33 +++++++++++++--------------------
 3 files changed, 18 insertions(+), 43 deletions(-)

diff --git a/include/subexpr.h b/include/subexpr.h
index 78b398c..8c5661b 100644
--- a/include/subexpr.h
+++ b/include/subexpr.h
@@ -57,11 +57,4 @@ typedef struct elementwise_mult_expr
     CSR_Matrix *CSR_work2;
 } elementwise_mult_expr;
 
-/* Promote (broadcast) to larger shape */
-typedef struct promote_expr
-{
-    expr base;
-    /* target shape stored in base.d1, base.d2 */
-} promote_expr;
-
 #endif /* SUBEXPR_H */
diff --git a/src/affine/neg.c b/src/affine/neg.c
index 6ea0831..632372b 100644
--- a/src/affine/neg.c
+++ b/src/affine/neg.c
@@ -1,4 +1,5 @@
 #include "affine.h"
+#include <string.h>
 
 static void forward(expr *node, const double *u)
 {
@@ -22,14 +23,8 @@ static void jacobian_init(expr *node)
     node->jacobian = new_csr_matrix(child_jac->m, child_jac->n, child_jac->nnz);
 
     /* copy row pointers and column indices (sparsity pattern is constant) */
-    for (int i = 0; i <= child_jac->m; i++)
-    {
-        node->jacobian->p[i] = child_jac->p[i];
-    }
-    for (int k = 0; k < child_jac->nnz; k++)
-    {
-        node->jacobian->i[k] = child_jac->i[k];
-    }
+    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;
 }
 
@@ -56,14 +51,8 @@ static void wsum_hess_init(expr *node)
     node->wsum_hess = new_csr_matrix(child_hess->m, child_hess->n, child_hess->nnz);
 
     /* copy row pointers and column indices (sparsity pattern is constant) */
-    for (int i = 0; i <= child_hess->m; i++)
-    {
-        node->wsum_hess->p[i] = child_hess->p[i];
-    }
-    for (int k = 0; k < child_hess->nnz; k++)
-    {
-        node->wsum_hess->i[k] = child_hess->i[k];
-    }
+    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;
 }
 
diff --git a/src/affine/promote.c b/src/affine/promote.c
index a40b7c9..a59b452 100644
--- a/src/affine/promote.c
+++ b/src/affine/promote.c
@@ -1,14 +1,12 @@
 #include "affine.h"
-#include "subexpr.h"
 #include <stdlib.h>
+#include <string.h>
 
 /* Promote broadcasts a child expression to a larger shape.
  * Typically used to broadcast a scalar to a vector. */
 
 static void forward(expr *node, const double *u)
 {
-    promote_expr *prom = (promote_expr *) node;
-
     /* child's forward pass */
     node->left->forward(node->left, u);
 
@@ -19,7 +17,6 @@ static void forward(expr *node, const double *u)
         /* replicate pattern: output[i] = child[i % child_size] */
         node->value[i] = node->left->value[i % child_size];
     }
-    (void) prom; /* unused for now, shape info stored in d1/d2 */
 }
 
 static void jacobian_init(expr *node)
@@ -85,22 +82,20 @@ static void wsum_hess_init(expr *node)
     node->wsum_hess = new_csr_matrix(child_hess->m, child_hess->n, child_hess->nnz);
 
     /* copy row pointers and column indices (sparsity pattern is constant) */
-    for (int i = 0; i <= child_hess->m; i++)
-    {
-        node->wsum_hess->p[i] = child_hess->p[i];
-    }
-    for (int k = 0; k < child_hess->nnz; k++)
-    {
-        node->wsum_hess->i[k] = child_hess->i[k];
-    }
+    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;
+
+    /* allocate workspace for summing weights */
+    node->dwork = (double *) malloc(node->left->size * sizeof(double));
 }
 
 static void eval_wsum_hess(expr *node, const double *w)
 {
     /* Sum weights that correspond to the same child element */
     int child_size = node->left->size;
-    double *summed_w = (double *) calloc(child_size, sizeof(double));
+    double *summed_w = node->dwork;
+    memset(summed_w, 0, child_size * sizeof(double));
     for (int i = 0; i < node->size; i++)
     {
         summed_w[i % child_size] += w[i];
@@ -108,7 +103,6 @@ static void eval_wsum_hess(expr *node, const double *w)
 
     /* evaluate child's wsum_hess with summed weights */
     node->left->eval_wsum_hess(node->left, summed_w);
-    free(summed_w);
 
     /* copy values only (sparsity pattern set in wsum_hess_init) */
     CSR_Matrix *child_hess = node->left->wsum_hess;
@@ -125,12 +119,11 @@ static bool is_affine(const expr *node)
 
 expr *new_promote(expr *child, int d1, int d2)
 {
-    /* Allocate the type-specific struct */
-    promote_expr *prom = (promote_expr *) calloc(1, sizeof(promote_expr));
-    expr *node = &prom->base;
-
-    init_expr(node, d1, d2, child->n_vars, forward, jacobian_init, eval_jacobian,
-              is_affine, NULL);
+    expr *node = new_expr(d1, d2, child->n_vars);
+    node->forward = forward;
+    node->jacobian_init = jacobian_init;
+    node->eval_jacobian = eval_jacobian;
+    node->is_affine = is_affine;
 
     node->left = child;
     expr_retain(child);

From 4bb8351b191924a3d50505018cb2ea6a388517ec Mon Sep 17 00:00:00 2001
From: William Zijie Zhang <william@gridmatic.com>
Date: Wed, 7 Jan 2026 15:57:34 -0500
Subject: [PATCH 22/27] Use memcpy for value copies and simplify row pointer
 loop
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

- promote.c: Use memcpy for wsum_hess values in eval_wsum_hess
- problem.c: Simplify row pointer loop to directly copy cjac->p with offset

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
---
 src/affine/promote.c | 5 +----
 src/problem.c        | 5 ++---
 2 files changed, 3 insertions(+), 7 deletions(-)

diff --git a/src/affine/promote.c b/src/affine/promote.c
index a59b452..efebf9a 100644
--- a/src/affine/promote.c
+++ b/src/affine/promote.c
@@ -106,10 +106,7 @@ static void eval_wsum_hess(expr *node, const double *w)
 
     /* copy values only (sparsity pattern set in wsum_hess_init) */
     CSR_Matrix *child_hess = node->left->wsum_hess;
-    for (int k = 0; k < child_hess->nnz; k++)
-    {
-        node->wsum_hess->x[k] = child_hess->x[k];
-    }
+    memcpy(node->wsum_hess->x, child_hess->x, child_hess->nnz * sizeof(double));
 }
 
 static bool is_affine(const expr *node)
diff --git a/src/problem.c b/src/problem.c
index 6d9dec8..945d0a2 100644
--- a/src/problem.c
+++ b/src/problem.c
@@ -159,10 +159,9 @@ void problem_jacobian(problem *prob)
         CSR_Matrix *cjac = c->jacobian;
 
         /* Copy row pointers with offset */
-        for (int r = 0; r < cjac->m; r++)
+        for (int r = 1; r <= cjac->m; r++)
         {
-            int row_nnz = cjac->p[r + 1] - cjac->p[r];
-            stacked->p[row_offset + r + 1] = stacked->p[row_offset + r] + row_nnz;
+            stacked->p[row_offset + r] = nnz_offset + cjac->p[r];
         }
 
         /* Copy column indices and values */

From 856ca50eb7f28067e27b8ff1a350da894dcb0520 Mon Sep 17 00:00:00 2001
From: William Zijie Zhang <william@gridmatic.com>
Date: Thu, 8 Jan 2026 10:50:55 -0500
Subject: [PATCH 23/27] Simplify promote to only handle scalar inputs

This matches CVXPY's promote atom behavior which only broadcasts scalars
to vectors/matrices. Removes vector-to-vector promotion support and
the associated dwork allocation.

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
---
 src/affine/promote.c                     | 70 +++++++++---------------
 tests/all_tests.c                        |  2 +-
 tests/forward_pass/affine/test_promote.h | 43 +++++++--------
 3 files changed, 45 insertions(+), 70 deletions(-)

diff --git a/src/affine/promote.c b/src/affine/promote.c
index efebf9a..a9202c5 100644
--- a/src/affine/promote.c
+++ b/src/affine/promote.c
@@ -2,50 +2,41 @@
 #include <stdlib.h>
 #include <string.h>
 
-/* Promote broadcasts a child expression to a larger shape.
- * Typically used to broadcast a scalar to a vector. */
+/* Promote broadcasts a scalar expression to a vector/matrix shape.
+ * This matches CVXPY's promote atom which only handles scalars. */
 
 static void forward(expr *node, const double *u)
 {
-    /* child's forward pass */
     node->left->forward(node->left, u);
 
-    /* broadcast child value to output shape */
-    int child_size = node->left->size;
+    /* broadcast scalar value to all output elements */
+    double val = node->left->value[0];
     for (int i = 0; i < node->size; i++)
     {
-        /* replicate pattern: output[i] = child[i % child_size] */
-        node->value[i] = node->left->value[i % child_size];
+        node->value[i] = val;
     }
 }
 
 static void jacobian_init(expr *node)
 {
-    /* initialize child's jacobian */
     node->left->jacobian_init(node->left);
 
-    /* Each output row copies a row from child's jacobian (with wrapping).
-     * For scalar->vector: all rows are copies of the single child row.
-     * nnz = (output_size / child_size) * child_jac_nnz */
+    /* Each output row copies the single row from child's jacobian */
     CSR_Matrix *child_jac = node->left->jacobian;
-    int child_size = node->left->size;
-    int repeat = (node->size + child_size - 1) / child_size;
-    int nnz_max = repeat * child_jac->nnz;
-    if (nnz_max == 0) nnz_max = 1;
+    int nnz = node->size * child_jac->nnz;
+    if (nnz == 0) nnz = 1;
 
-    node->jacobian = new_csr_matrix(node->size, node->n_vars, nnz_max);
+    node->jacobian = new_csr_matrix(node->size, node->n_vars, nnz);
     CSR_Matrix *jac = node->jacobian;
 
-    /* Build sparsity pattern by replicating child's rows */
+    /* Build sparsity pattern by replicating child's single row */
     jac->nnz = 0;
     for (int row = 0; row < node->size; row++)
     {
         jac->p[row] = jac->nnz;
-        int child_row = row % child_size;
-        for (int k = child_jac->p[child_row]; k < child_jac->p[child_row + 1]; k++)
+        for (int k = child_jac->p[0]; k < child_jac->p[1]; k++)
         {
-            jac->i[jac->nnz] = child_jac->i[k];
-            jac->nnz++;
+            jac->i[jac->nnz++] = child_jac->i[k];
         }
     }
     jac->p[node->size] = jac->nnz;
@@ -53,58 +44,47 @@ static void jacobian_init(expr *node)
 
 static void eval_jacobian(expr *node)
 {
-    /* evaluate child's jacobian */
     node->left->eval_jacobian(node->left);
 
     CSR_Matrix *child_jac = node->left->jacobian;
     CSR_Matrix *jac = node->jacobian;
-    int child_size = node->left->size;
+    int child_nnz = child_jac->p[1] - child_jac->p[0];
 
-    /* Copy values only (sparsity pattern set in jacobian_init) */
-    int idx = 0;
+    /* Copy child's row values to each output row */
     for (int row = 0; row < node->size; row++)
     {
-        int child_row = row % child_size;
-        for (int k = child_jac->p[child_row]; k < child_jac->p[child_row + 1]; k++)
-        {
-            jac->x[idx++] = child_jac->x[k];
-        }
+        memcpy(&jac->x[row * child_nnz], &child_jac->x[child_jac->p[0]],
+               child_nnz * sizeof(double));
     }
 }
 
 static void wsum_hess_init(expr *node)
 {
-    /* initialize child's wsum_hess */
     node->left->wsum_hess_init(node->left);
 
     /* same sparsity as child since we're summing weights */
     CSR_Matrix *child_hess = node->left->wsum_hess;
     node->wsum_hess = new_csr_matrix(child_hess->m, child_hess->n, child_hess->nnz);
 
-    /* copy row pointers and column indices (sparsity pattern is constant) */
+    /* copy sparsity pattern */
     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;
-
-    /* allocate workspace for summing weights */
-    node->dwork = (double *) malloc(node->left->size * sizeof(double));
 }
 
 static void eval_wsum_hess(expr *node, const double *w)
 {
-    /* Sum weights that correspond to the same child element */
-    int child_size = node->left->size;
-    double *summed_w = node->dwork;
-    memset(summed_w, 0, child_size * sizeof(double));
+    /* Sum all weights (they all correspond to the same scalar child) */
+    double sum_w = 0.0;
     for (int i = 0; i < node->size; i++)
     {
-        summed_w[i % child_size] += w[i];
+        sum_w += w[i];
     }
 
-    /* evaluate child's wsum_hess with summed weights */
-    node->left->eval_wsum_hess(node->left, summed_w);
+    /* evaluate child's wsum_hess with summed weight */
+    node->left->eval_wsum_hess(node->left, &sum_w);
 
-    /* copy values only (sparsity pattern set in wsum_hess_init) */
+    /* copy values */
     CSR_Matrix *child_hess = node->left->wsum_hess;
     memcpy(node->wsum_hess->x, child_hess->x, child_hess->nnz * sizeof(double));
 }
@@ -121,11 +101,11 @@ expr *new_promote(expr *child, int d1, int d2)
     node->jacobian_init = jacobian_init;
     node->eval_jacobian = eval_jacobian;
     node->is_affine = is_affine;
+    node->wsum_hess_init = wsum_hess_init;
+    node->eval_wsum_hess = eval_wsum_hess;
 
     node->left = child;
     expr_retain(child);
-    node->wsum_hess_init = wsum_hess_init;
-    node->eval_wsum_hess = eval_wsum_hess;
 
     return node;
 }
diff --git a/tests/all_tests.c b/tests/all_tests.c
index 1994876..19ac141 100644
--- a/tests/all_tests.c
+++ b/tests/all_tests.c
@@ -55,7 +55,7 @@ int main(void)
     mu_run_test(test_neg_chain, tests_run);
     mu_run_test(test_promote_scalar_to_vector, tests_run);
     mu_run_test(test_promote_scalar_jacobian, tests_run);
-    mu_run_test(test_promote_vector_jacobian, tests_run);
+    mu_run_test(test_promote_scalar_to_matrix_jacobian, tests_run);
     mu_run_test(test_exp, tests_run);
     mu_run_test(test_log, tests_run);
     mu_run_test(test_composite, tests_run);
diff --git a/tests/forward_pass/affine/test_promote.h b/tests/forward_pass/affine/test_promote.h
index 7195fb0..dd47c39 100644
--- a/tests/forward_pass/affine/test_promote.h
+++ b/tests/forward_pass/affine/test_promote.h
@@ -49,38 +49,33 @@ const char *test_promote_scalar_jacobian(void)
     return 0;
 }
 
-const char *test_promote_vector_jacobian(void)
+const char *test_promote_scalar_to_matrix_jacobian(void)
 {
-    /* Promote 2-vector to 4-vector, check jacobian */
-    double u[2] = {1.0, 2.0};
-    expr *var = new_variable(2, 1, 0, 2);
-    expr *promote_node = new_promote(var, 4, 1);
+    /* Promote scalar to 2x3 matrix, check jacobian */
+    double u[1] = {7.0};
+    expr *var = new_variable(1, 1, 0, 1);
+    expr *promote_node = new_promote(var, 2, 3);
     promote_node->forward(promote_node, u);
 
-    /* Pattern repeats: [1, 2, 1, 2] */
-    double expected_val[4] = {1.0, 2.0, 1.0, 2.0};
-    mu_assert("promote vector forward failed",
-              cmp_double_array(promote_node->value, expected_val, 4));
+    /* Forward: all 6 elements should be 7.0 */
+    double expected_val[6] = {7.0, 7.0, 7.0, 7.0, 7.0, 7.0};
+    mu_assert("promote scalar->matrix forward failed",
+              cmp_double_array(promote_node->value, expected_val, 6));
 
     promote_node->jacobian_init(promote_node);
     promote_node->eval_jacobian(promote_node);
 
-    /* Jacobian is 4x2:
-     * Row 0: [1, 0] (output 0 depends on input 0)
-     * Row 1: [0, 1] (output 1 depends on input 1)
-     * Row 2: [1, 0] (output 2 depends on input 0)
-     * Row 3: [0, 1] (output 3 depends on input 1)
-     */
-    double expected_x[4] = {1.0, 1.0, 1.0, 1.0};
-    int expected_p[5] = {0, 1, 2, 3, 4};
-    int expected_i[4] = {0, 1, 0, 1};
+    /* Jacobian is 6x1 with all 1s (each output depends on same scalar input) */
+    double expected_x[6] = {1.0, 1.0, 1.0, 1.0, 1.0, 1.0};
+    int expected_p[7] = {0, 1, 2, 3, 4, 5, 6};
+    int expected_i[6] = {0, 0, 0, 0, 0, 0};
 
-    mu_assert("promote vector jacobian vals fail",
-              cmp_double_array(promote_node->jacobian->x, expected_x, 4));
-    mu_assert("promote vector jacobian rows fail",
-              cmp_int_array(promote_node->jacobian->p, expected_p, 5));
-    mu_assert("promote vector jacobian cols fail",
-              cmp_int_array(promote_node->jacobian->i, expected_i, 4));
+    mu_assert("promote matrix jacobian vals fail",
+              cmp_double_array(promote_node->jacobian->x, expected_x, 6));
+    mu_assert("promote matrix jacobian rows fail",
+              cmp_int_array(promote_node->jacobian->p, expected_p, 7));
+    mu_assert("promote matrix jacobian cols fail",
+              cmp_int_array(promote_node->jacobian->i, expected_i, 6));
 
     free_expr(promote_node);
     return 0;

From 012efc6f25064d9aa50f405ec9cc1c7acc3c4809 Mon Sep 17 00:00:00 2001
From: William Zijie Zhang <william@gridmatic.com>
Date: Thu, 8 Jan 2026 14:16:51 -0500
Subject: [PATCH 24/27] Refactor promote: use memcpy and reorganize tests

- Use memcpy instead of loops for copying jacobian indices and values
- Use pointer arithmetic style (ptr + offset) instead of &ptr[offset]
- Remove unnecessary nnz == 0 guard in jacobian_init
- Move jacobian tests from forward_pass/ to jacobian_tests/

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
---
 src/affine/promote.c                     | 11 ++--
 tests/all_tests.c                        |  5 +-
 tests/forward_pass/affine/test_promote.h | 58 ---------------------
 tests/jacobian_tests/test_promote.h      | 66 ++++++++++++++++++++++++
 4 files changed, 74 insertions(+), 66 deletions(-)
 create mode 100644 tests/jacobian_tests/test_promote.h

diff --git a/src/affine/promote.c b/src/affine/promote.c
index a9202c5..fca5c80 100644
--- a/src/affine/promote.c
+++ b/src/affine/promote.c
@@ -24,20 +24,19 @@ static void jacobian_init(expr *node)
     /* Each output row copies the single row from child's jacobian */
     CSR_Matrix *child_jac = node->left->jacobian;
     int nnz = node->size * child_jac->nnz;
-    if (nnz == 0) nnz = 1;
 
     node->jacobian = new_csr_matrix(node->size, node->n_vars, nnz);
     CSR_Matrix *jac = node->jacobian;
 
     /* Build sparsity pattern by replicating child's single row */
+    int child_nnz = child_jac->p[1] - child_jac->p[0];
     jac->nnz = 0;
     for (int row = 0; row < node->size; row++)
     {
         jac->p[row] = jac->nnz;
-        for (int k = child_jac->p[0]; k < child_jac->p[1]; k++)
-        {
-            jac->i[jac->nnz++] = child_jac->i[k];
-        }
+        memcpy(jac->i + jac->nnz, child_jac->i + child_jac->p[0],
+               child_nnz * sizeof(int));
+        jac->nnz += child_nnz;
     }
     jac->p[node->size] = jac->nnz;
 }
@@ -53,7 +52,7 @@ static void eval_jacobian(expr *node)
     /* Copy child's row values to each output row */
     for (int row = 0; row < node->size; row++)
     {
-        memcpy(&jac->x[row * child_nnz], &child_jac->x[child_jac->p[0]],
+        memcpy(jac->x + row * child_nnz, child_jac->x + child_jac->p[0],
                child_nnz * sizeof(double));
     }
 }
diff --git a/tests/all_tests.c b/tests/all_tests.c
index 19ac141..01e82bb 100644
--- a/tests/all_tests.c
+++ b/tests/all_tests.c
@@ -17,6 +17,7 @@
 #include "jacobian_tests/test_elementwise_mult.h"
 #include "jacobian_tests/test_hstack.h"
 #include "jacobian_tests/test_log.h"
+#include "jacobian_tests/test_promote.h"
 #include "jacobian_tests/test_quad_form.h"
 #include "jacobian_tests/test_quad_over_lin.h"
 #include "jacobian_tests/test_rel_entr.h"
@@ -54,8 +55,6 @@ int main(void)
     mu_run_test(test_neg_jacobian, tests_run);
     mu_run_test(test_neg_chain, tests_run);
     mu_run_test(test_promote_scalar_to_vector, tests_run);
-    mu_run_test(test_promote_scalar_jacobian, tests_run);
-    mu_run_test(test_promote_scalar_to_matrix_jacobian, tests_run);
     mu_run_test(test_exp, tests_run);
     mu_run_test(test_log, tests_run);
     mu_run_test(test_composite, tests_run);
@@ -91,6 +90,8 @@ int main(void)
     mu_run_test(test_jacobian_sum_log_axis_1, tests_run);
     mu_run_test(test_jacobian_hstack_vectors, tests_run);
     mu_run_test(test_jacobian_hstack_matrix, tests_run);
+    mu_run_test(test_promote_scalar_jacobian, tests_run);
+    mu_run_test(test_promote_scalar_to_matrix_jacobian, tests_run);
     mu_run_test(test_wsum_hess_multiply_1, tests_run);
     mu_run_test(test_wsum_hess_multiply_2, tests_run);
 
diff --git a/tests/forward_pass/affine/test_promote.h b/tests/forward_pass/affine/test_promote.h
index dd47c39..dcc9060 100644
--- a/tests/forward_pass/affine/test_promote.h
+++ b/tests/forward_pass/affine/test_promote.h
@@ -22,61 +22,3 @@ const char *test_promote_scalar_to_vector(void)
     free_expr(promote_node);
     return 0;
 }
-
-const char *test_promote_scalar_jacobian(void)
-{
-    /* Promote scalar to 3-element vector, check jacobian */
-    double u[1] = {2.0};
-    expr *var = new_variable(1, 1, 0, 1);
-    expr *promote_node = new_promote(var, 3, 1);
-    promote_node->forward(promote_node, u);
-    promote_node->jacobian_init(promote_node);
-    promote_node->eval_jacobian(promote_node);
-
-    /* Jacobian is 3x1 with all 1s (each output depends on same input) */
-    double expected_x[3] = {1.0, 1.0, 1.0};
-    int expected_p[4] = {0, 1, 2, 3};
-    int expected_i[3] = {0, 0, 0};
-
-    mu_assert("promote jacobian vals fail",
-              cmp_double_array(promote_node->jacobian->x, expected_x, 3));
-    mu_assert("promote jacobian rows fail",
-              cmp_int_array(promote_node->jacobian->p, expected_p, 4));
-    mu_assert("promote jacobian cols fail",
-              cmp_int_array(promote_node->jacobian->i, expected_i, 3));
-
-    free_expr(promote_node);
-    return 0;
-}
-
-const char *test_promote_scalar_to_matrix_jacobian(void)
-{
-    /* Promote scalar to 2x3 matrix, check jacobian */
-    double u[1] = {7.0};
-    expr *var = new_variable(1, 1, 0, 1);
-    expr *promote_node = new_promote(var, 2, 3);
-    promote_node->forward(promote_node, u);
-
-    /* Forward: all 6 elements should be 7.0 */
-    double expected_val[6] = {7.0, 7.0, 7.0, 7.0, 7.0, 7.0};
-    mu_assert("promote scalar->matrix forward failed",
-              cmp_double_array(promote_node->value, expected_val, 6));
-
-    promote_node->jacobian_init(promote_node);
-    promote_node->eval_jacobian(promote_node);
-
-    /* Jacobian is 6x1 with all 1s (each output depends on same scalar input) */
-    double expected_x[6] = {1.0, 1.0, 1.0, 1.0, 1.0, 1.0};
-    int expected_p[7] = {0, 1, 2, 3, 4, 5, 6};
-    int expected_i[6] = {0, 0, 0, 0, 0, 0};
-
-    mu_assert("promote matrix jacobian vals fail",
-              cmp_double_array(promote_node->jacobian->x, expected_x, 6));
-    mu_assert("promote matrix jacobian rows fail",
-              cmp_int_array(promote_node->jacobian->p, expected_p, 7));
-    mu_assert("promote matrix jacobian cols fail",
-              cmp_int_array(promote_node->jacobian->i, expected_i, 6));
-
-    free_expr(promote_node);
-    return 0;
-}
diff --git a/tests/jacobian_tests/test_promote.h b/tests/jacobian_tests/test_promote.h
new file mode 100644
index 0000000..978d31b
--- /dev/null
+++ b/tests/jacobian_tests/test_promote.h
@@ -0,0 +1,66 @@
+#include <math.h>
+#include <stdio.h>
+#include <stdlib.h>
+
+#include "affine.h"
+#include "expr.h"
+#include "minunit.h"
+#include "test_helpers.h"
+
+const char *test_promote_scalar_jacobian(void)
+{
+    /* Promote scalar to 3-element vector, check jacobian */
+    double u[1] = {2.0};
+    expr *var = new_variable(1, 1, 0, 1);
+    expr *promote_node = new_promote(var, 3, 1);
+    promote_node->forward(promote_node, u);
+    promote_node->jacobian_init(promote_node);
+    promote_node->eval_jacobian(promote_node);
+
+    /* Jacobian is 3x1 with all 1s (each output depends on same input) */
+    double expected_x[3] = {1.0, 1.0, 1.0};
+    int expected_p[4] = {0, 1, 2, 3};
+    int expected_i[3] = {0, 0, 0};
+
+    mu_assert("promote jacobian vals fail",
+              cmp_double_array(promote_node->jacobian->x, expected_x, 3));
+    mu_assert("promote jacobian rows fail",
+              cmp_int_array(promote_node->jacobian->p, expected_p, 4));
+    mu_assert("promote jacobian cols fail",
+              cmp_int_array(promote_node->jacobian->i, expected_i, 3));
+
+    free_expr(promote_node);
+    return 0;
+}
+
+const char *test_promote_scalar_to_matrix_jacobian(void)
+{
+    /* Promote scalar to 2x3 matrix, check jacobian */
+    double u[1] = {7.0};
+    expr *var = new_variable(1, 1, 0, 1);
+    expr *promote_node = new_promote(var, 2, 3);
+    promote_node->forward(promote_node, u);
+
+    /* Forward: all 6 elements should be 7.0 */
+    double expected_val[6] = {7.0, 7.0, 7.0, 7.0, 7.0, 7.0};
+    mu_assert("promote scalar->matrix forward failed",
+              cmp_double_array(promote_node->value, expected_val, 6));
+
+    promote_node->jacobian_init(promote_node);
+    promote_node->eval_jacobian(promote_node);
+
+    /* Jacobian is 6x1 with all 1s (each output depends on same scalar input) */
+    double expected_x[6] = {1.0, 1.0, 1.0, 1.0, 1.0, 1.0};
+    int expected_p[7] = {0, 1, 2, 3, 4, 5, 6};
+    int expected_i[6] = {0, 0, 0, 0, 0, 0};
+
+    mu_assert("promote matrix jacobian vals fail",
+              cmp_double_array(promote_node->jacobian->x, expected_x, 6));
+    mu_assert("promote matrix jacobian rows fail",
+              cmp_int_array(promote_node->jacobian->p, expected_p, 7));
+    mu_assert("promote matrix jacobian cols fail",
+              cmp_int_array(promote_node->jacobian->i, expected_i, 6));
+
+    free_expr(promote_node);
+    return 0;
+}

From 40927d5a6bc824df5a6635d7c9bec2a31e13efee Mon Sep 17 00:00:00 2001
From: William Zijie Zhang <william@gridmatic.com>
Date: Fri, 9 Jan 2026 12:45:20 -0500
Subject: [PATCH 25/27] Fix duplicate eval_jacobian definition in variable.c

Remove duplicate function definition that was causing compilation errors,
likely introduced during merge from main.

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
---
 src/affine/variable.c | 6 ------
 1 file changed, 6 deletions(-)

diff --git a/src/affine/variable.c b/src/affine/variable.c
index 9d9f263..8ce8c6e 100644
--- a/src/affine/variable.c
+++ b/src/affine/variable.c
@@ -19,12 +19,6 @@ static void jacobian_init(expr *node)
     node->jacobian->p[node->size] = node->size;
 }
 
-static void eval_jacobian(expr *node)
-{
-    /* Jacobian is already initialized with correct values */
-    (void) node;
-}
-
 static void eval_jacobian(expr *node)
 {
     /* Variable jacobian never changes - nothing to evaluate */

From 78fd177b169824fe09f06abdc0990b431554a1d2 Mon Sep 17 00:00:00 2001
From: William Zijie Zhang <william@gridmatic.com>
Date: Fri, 9 Jan 2026 12:57:22 -0500
Subject: [PATCH 26/27] Refactor problem tests to use cmp_double_array and
 cmp_int_array

Replace manual element-by-element array assertions with helper functions
for cleaner, more maintainable test code.

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
---
 tests/problem/test_problem.h | 59 ++++++++++++++----------------------
 1 file changed, 22 insertions(+), 37 deletions(-)

diff --git a/tests/problem/test_problem.h b/tests/problem/test_problem.h
index 1c0e96b..7ea502b 100644
--- a/tests/problem/test_problem.h
+++ b/tests/problem/test_problem.h
@@ -71,9 +71,9 @@ const char *test_problem_objective_forward(void)
     problem_constraint_forward(prob, u);
 
     /* Constraint values should be [1, 2, 3] */
-    mu_assert("constraint[0] wrong", fabs(prob->constraint_values[0] - 1.0) < 1e-10);
-    mu_assert("constraint[1] wrong", fabs(prob->constraint_values[1] - 2.0) < 1e-10);
-    mu_assert("constraint[2] wrong", fabs(prob->constraint_values[2] - 3.0) < 1e-10);
+    double expected_constraints[3] = {1.0, 2.0, 3.0};
+    mu_assert("constraint values wrong",
+              cmp_double_array(prob->constraint_values, expected_constraints, 3));
 
     free_problem(prob);
 
@@ -111,17 +111,10 @@ const char *test_problem_constraint_forward(void)
 
     problem_constraint_forward(prob, u);
 
-    /* Check constraint values:
-     * [log(2), log(4), exp(2), exp(4)]
-     */
-    mu_assert("constraint[0] wrong",
-              fabs(prob->constraint_values[0] - log(2.0)) < 1e-10);
-    mu_assert("constraint[1] wrong",
-              fabs(prob->constraint_values[1] - log(4.0)) < 1e-10);
-    mu_assert("constraint[2] wrong",
-              fabs(prob->constraint_values[2] - exp(2.0)) < 1e-10);
-    mu_assert("constraint[3] wrong",
-              fabs(prob->constraint_values[3] - exp(4.0)) < 1e-10);
+    /* Check constraint values: [log(2), log(4), exp(2), exp(4)] */
+    double expected_constraints[4] = {log(2.0), log(4.0), exp(2.0), exp(4.0)};
+    mu_assert("constraint values wrong",
+              cmp_double_array(prob->constraint_values, expected_constraints, 4));
 
     free_problem(prob);
 
@@ -146,9 +139,9 @@ const char *test_problem_gradient(void)
     problem_gradient(prob);
 
     /* Expected gradient: [1/1, 1/2, 1/4] = [1.0, 0.5, 0.25] */
-    mu_assert("grad[0] wrong", fabs(prob->gradient_values[0] - 1.0) < 1e-10);
-    mu_assert("grad[1] wrong", fabs(prob->gradient_values[1] - 0.5) < 1e-10);
-    mu_assert("grad[2] wrong", fabs(prob->gradient_values[2] - 0.25) < 1e-10);
+    double expected_grad[3] = {1.0, 0.5, 0.25};
+    mu_assert("gradient wrong",
+              cmp_double_array(prob->gradient_values, expected_grad, 3));
 
     free_problem(prob);
 
@@ -188,17 +181,16 @@ const char *test_problem_jacobian(void)
     mu_assert("jac cols wrong", jac->n == 2);
 
     /* Check row pointers: each row has 1 element */
-    mu_assert("jac->p[0] wrong", jac->p[0] == 0);
-    mu_assert("jac->p[1] wrong", jac->p[1] == 1);
-    mu_assert("jac->p[2] wrong", jac->p[2] == 2);
+    int expected_p[3] = {0, 1, 2};
+    mu_assert("jac->p wrong", cmp_int_array(jac->p, expected_p, 3));
 
     /* Check column indices */
-    mu_assert("jac->i[0] wrong", jac->i[0] == 0);
-    mu_assert("jac->i[1] wrong", jac->i[1] == 1);
+    int expected_i[2] = {0, 1};
+    mu_assert("jac->i wrong", cmp_int_array(jac->i, expected_i, 2));
 
     /* Check values: [1/2, 1/4] */
-    mu_assert("jac->x[0] wrong", fabs(jac->x[0] - 0.5) < 1e-10);
-    mu_assert("jac->x[1] wrong", fabs(jac->x[1] - 0.25) < 1e-10);
+    double expected_x[2] = {0.5, 0.25};
+    mu_assert("jac->x wrong", cmp_double_array(jac->x, expected_x, 2));
 
     free_problem(prob);
 
@@ -254,17 +246,12 @@ const char *test_problem_jacobian_multi(void)
     mu_assert("jac nnz wrong", jac->nnz == 4);
 
     /* Check row pointers: each row has 1 element */
-    mu_assert("jac->p[0] wrong", jac->p[0] == 0);
-    mu_assert("jac->p[1] wrong", jac->p[1] == 1);
-    mu_assert("jac->p[2] wrong", jac->p[2] == 2);
-    mu_assert("jac->p[3] wrong", jac->p[3] == 3);
-    mu_assert("jac->p[4] wrong", jac->p[4] == 4);
+    int expected_p[5] = {0, 1, 2, 3, 4};
+    mu_assert("jac->p wrong", cmp_int_array(jac->p, expected_p, 5));
 
     /* Check column indices: diagonal pattern */
-    mu_assert("jac->i[0] wrong", jac->i[0] == 0);
-    mu_assert("jac->i[1] wrong", jac->i[1] == 1);
-    mu_assert("jac->i[2] wrong", jac->i[2] == 0);
-    mu_assert("jac->i[3] wrong", jac->i[3] == 1);
+    int expected_i[4] = {0, 1, 0, 1};
+    mu_assert("jac->i wrong", cmp_int_array(jac->i, expected_i, 4));
 
     /* Check values:
      * Row 0: 1/2 = 0.5
@@ -272,10 +259,8 @@ const char *test_problem_jacobian_multi(void)
      * Row 2: exp(2) ≈ 7.389
      * Row 3: exp(4) ≈ 54.598
      */
-    mu_assert("jac->x[0] wrong", fabs(jac->x[0] - 0.5) < 1e-10);
-    mu_assert("jac->x[1] wrong", fabs(jac->x[1] - 0.25) < 1e-10);
-    mu_assert("jac->x[2] wrong", fabs(jac->x[2] - exp(2.0)) < 1e-10);
-    mu_assert("jac->x[3] wrong", fabs(jac->x[3] - exp(4.0)) < 1e-10);
+    double expected_x[4] = {0.5, 0.25, exp(2.0), exp(4.0)};
+    mu_assert("jac->x wrong", cmp_double_array(jac->x, expected_x, 4));
 
     free_problem(prob);
 

From 55ca55f127e34abb889723f330b293ddc432f14e Mon Sep 17 00:00:00 2001
From: William Zijie Zhang <william@gridmatic.com>
Date: Fri, 9 Jan 2026 12:59:04 -0500
Subject: [PATCH 27/27] Move neg jacobian tests to
 tests/jacobian_tests/test_neg.h

Relocate test_neg_jacobian and test_neg_chain from forward_pass/affine/
to jacobian_tests/ for better test organization.

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
---
 tests/all_tests.c                    |  5 ++-
 tests/forward_pass/affine/test_neg.h | 56 -------------------------
 tests/jacobian_tests/test_neg.h      | 62 ++++++++++++++++++++++++++++
 3 files changed, 65 insertions(+), 58 deletions(-)
 create mode 100644 tests/jacobian_tests/test_neg.h

diff --git a/tests/all_tests.c b/tests/all_tests.c
index 4959659..502116e 100644
--- a/tests/all_tests.c
+++ b/tests/all_tests.c
@@ -15,6 +15,7 @@
 #include "forward_pass/elementwise/test_log.h"
 #include "jacobian_tests/test_composite.h"
 #include "jacobian_tests/test_elementwise_mult.h"
+#include "jacobian_tests/test_neg.h"
 #include "jacobian_tests/test_hstack.h"
 #include "jacobian_tests/test_log.h"
 #include "jacobian_tests/test_promote.h"
@@ -56,8 +57,6 @@ int main(void)
     mu_run_test(test_addition, tests_run);
     mu_run_test(test_linear_op, tests_run);
     mu_run_test(test_neg_forward, tests_run);
-    mu_run_test(test_neg_jacobian, tests_run);
-    mu_run_test(test_neg_chain, tests_run);
     mu_run_test(test_promote_scalar_to_vector, tests_run);
     mu_run_test(test_exp, tests_run);
     mu_run_test(test_log, tests_run);
@@ -69,6 +68,8 @@ int main(void)
     mu_run_test(test_hstack_forward_matrix, tests_run);
 
     printf("\n--- Jacobian Tests ---\n");
+    mu_run_test(test_neg_jacobian, tests_run);
+    mu_run_test(test_neg_chain, tests_run);
     mu_run_test(test_jacobian_log, tests_run);
     mu_run_test(test_jacobian_log_matrix, tests_run);
     mu_run_test(test_jacobian_composite_log, tests_run);
diff --git a/tests/forward_pass/affine/test_neg.h b/tests/forward_pass/affine/test_neg.h
index af0c9aa..0fad2b1 100644
--- a/tests/forward_pass/affine/test_neg.h
+++ b/tests/forward_pass/affine/test_neg.h
@@ -18,59 +18,3 @@ const char *test_neg_forward(void)
     free_expr(neg_node);
     return 0;
 }
-
-const char *test_neg_jacobian(void)
-{
-    double u[3] = {1.0, 2.0, 3.0};
-    expr *var = new_variable(3, 1, 0, 3);
-    expr *neg_node = new_neg(var);
-    neg_node->forward(neg_node, u);
-    neg_node->jacobian_init(neg_node);
-    neg_node->eval_jacobian(neg_node);
-
-    /* Jacobian of neg(x) is -I (diagonal with -1) */
-    double expected_x[3] = {-1.0, -1.0, -1.0};
-    int expected_p[4] = {0, 1, 2, 3};
-    int expected_i[3] = {0, 1, 2};
-
-    mu_assert("neg jacobian vals fail",
-              cmp_double_array(neg_node->jacobian->x, expected_x, 3));
-    mu_assert("neg jacobian rows fail",
-              cmp_int_array(neg_node->jacobian->p, expected_p, 4));
-    mu_assert("neg jacobian cols fail",
-              cmp_int_array(neg_node->jacobian->i, expected_i, 3));
-
-    free_expr(neg_node);
-    return 0;
-}
-
-const char *test_neg_chain(void)
-{
-    /* Test neg(neg(x)) = x */
-    double u[3] = {1.0, 2.0, 3.0};
-    expr *var = new_variable(3, 1, 0, 3);
-    expr *neg1 = new_neg(var);
-    expr *neg2 = new_neg(neg1);
-    neg2->forward(neg2, u);
-
-    /* neg(neg(x)) should equal x */
-    mu_assert("neg chain forward failed", cmp_double_array(neg2->value, u, 3));
-
-    neg2->jacobian_init(neg2);
-    neg2->eval_jacobian(neg2);
-
-    /* Jacobian of neg(neg(x)) is (-1)*(-1)*I = I */
-    double expected_x[3] = {1.0, 1.0, 1.0};
-    int expected_p[4] = {0, 1, 2, 3};
-    int expected_i[3] = {0, 1, 2};
-
-    mu_assert("neg chain jacobian vals fail",
-              cmp_double_array(neg2->jacobian->x, expected_x, 3));
-    mu_assert("neg chain jacobian rows fail",
-              cmp_int_array(neg2->jacobian->p, expected_p, 4));
-    mu_assert("neg chain jacobian cols fail",
-              cmp_int_array(neg2->jacobian->i, expected_i, 3));
-
-    free_expr(neg2);
-    return 0;
-}
diff --git a/tests/jacobian_tests/test_neg.h b/tests/jacobian_tests/test_neg.h
new file mode 100644
index 0000000..f5377dc
--- /dev/null
+++ b/tests/jacobian_tests/test_neg.h
@@ -0,0 +1,62 @@
+#include <stdio.h>
+
+#include "affine.h"
+#include "expr.h"
+#include "minunit.h"
+#include "test_helpers.h"
+
+const char *test_neg_jacobian(void)
+{
+    double u[3] = {1.0, 2.0, 3.0};
+    expr *var = new_variable(3, 1, 0, 3);
+    expr *neg_node = new_neg(var);
+    neg_node->forward(neg_node, u);
+    neg_node->jacobian_init(neg_node);
+    neg_node->eval_jacobian(neg_node);
+
+    /* Jacobian of neg(x) is -I (diagonal with -1) */
+    double expected_x[3] = {-1.0, -1.0, -1.0};
+    int expected_p[4] = {0, 1, 2, 3};
+    int expected_i[3] = {0, 1, 2};
+
+    mu_assert("neg jacobian vals fail",
+              cmp_double_array(neg_node->jacobian->x, expected_x, 3));
+    mu_assert("neg jacobian rows fail",
+              cmp_int_array(neg_node->jacobian->p, expected_p, 4));
+    mu_assert("neg jacobian cols fail",
+              cmp_int_array(neg_node->jacobian->i, expected_i, 3));
+
+    free_expr(neg_node);
+    return 0;
+}
+
+const char *test_neg_chain(void)
+{
+    /* Test neg(neg(x)) = x */
+    double u[3] = {1.0, 2.0, 3.0};
+    expr *var = new_variable(3, 1, 0, 3);
+    expr *neg1 = new_neg(var);
+    expr *neg2 = new_neg(neg1);
+    neg2->forward(neg2, u);
+
+    /* neg(neg(x)) should equal x */
+    mu_assert("neg chain forward failed", cmp_double_array(neg2->value, u, 3));
+
+    neg2->jacobian_init(neg2);
+    neg2->eval_jacobian(neg2);
+
+    /* Jacobian of neg(neg(x)) is (-1)*(-1)*I = I */
+    double expected_x[3] = {1.0, 1.0, 1.0};
+    int expected_p[4] = {0, 1, 2, 3};
+    int expected_i[3] = {0, 1, 2};
+
+    mu_assert("neg chain jacobian vals fail",
+              cmp_double_array(neg2->jacobian->x, expected_x, 3));
+    mu_assert("neg chain jacobian rows fail",
+              cmp_int_array(neg2->jacobian->p, expected_p, 4));
+    mu_assert("neg chain jacobian cols fail",
+              cmp_int_array(neg2->jacobian->i, expected_i, 3));
+
+    free_expr(neg2);
+    return 0;
+}