Prep for CoeffExtractor

include/problem.h

@@ -26,6 +26,7 @@ typedef struct problem
 /* Retains objective and constraints (shared ownership with caller) */
 problem *new_problem(expr *objective, expr **constraints, int n_constraints);
 void problem_init_derivatives(problem *prob);
+void problem_init_jacobian_only(problem *prob);  /* Skip Hessian setup */
 void free_problem(problem *prob);
 
 double problem_objective_forward(problem *prob, const double *u);

python/bindings.c

@@ -38,6 +38,7 @@
 #include "problem/gradient.h"
 #include "problem/hessian.h"
 #include "problem/init_derivatives.h"
+#include "problem/init_jacobian_only.h"
 #include "problem/jacobian.h"
 #include "problem/make_problem.h"
 #include "problem/objective_forward.h"
@@ -94,6 +95,8 @@ static PyMethodDef DNLPMethods[] = {
      "Create problem from objective and constraints"},
     {"problem_init_derivatives", py_problem_init_derivatives, METH_VARARGS,
      "Initialize derivative structures"},
+    {"problem_init_jacobian_only", py_problem_init_jacobian_only, METH_VARARGS,
+     "Initialize Jacobian structure only (skip Hessian)"},
     {"problem_objective_forward", py_problem_objective_forward, METH_VARARGS,
      "Evaluate objective only"},
     {"problem_constraint_forward", py_problem_constraint_forward, METH_VARARGS,

python/problem/init_jacobian_only.h

@@ -0,0 +1,27 @@
+#ifndef PROBLEM_INIT_JACOBIAN_ONLY_H
+#define PROBLEM_INIT_JACOBIAN_ONLY_H
+
+#include "common.h"
+
+static PyObject *py_problem_init_jacobian_only(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_jacobian_only(prob);
+
+    Py_RETURN_NONE;
+}
+
+#endif /* PROBLEM_INIT_JACOBIAN_ONLY_H */

src/bivariate/left_matmul.c

@@ -64,11 +64,37 @@ static void jacobian_init(expr *node)
     expr *x = node->left;
     left_matmul_expr *lin_node = (left_matmul_expr *) node;
 
-    /* initialize child's jacobian and precompute sparsity of its transpose */
+    /* initialize child's jacobian */
     x->jacobian_init(x);
-    lin_node->CSC_work = csr_to_csc_fill_sparsity(x->jacobian, node->iwork);
 
-    /* precompute sparsity of this node's jacobian */
+    /* Fast path: child is a variable with identity Jacobian.
+     * For y = A @ x, J_y = A @ J_x where J_x is identity shifted by var_id.
+     * So J_y is just A with column indices shifted by var_id.
+     * This avoids O(nnz(A) * n) CSR-CSC multiplication. */
+    if (x->var_id != NOT_A_VARIABLE)
+    {
+        CSR_Matrix *A = lin_node->A;
+        node->jacobian = new_csr_matrix(A->m, node->n_vars, A->nnz);
+
+        /* Copy row pointers directly */
+        memcpy(node->jacobian->p, A->p, (A->m + 1) * sizeof(int));
+
+        /* Copy column indices with offset by var_id */
+        for (int k = 0; k < A->nnz; k++)
+        {
+            node->jacobian->i[k] = A->i[k] + x->var_id;
+        }
+
+        /* Copy values directly (A is constant, variable Jacobian entries are 1.0) */
+        memcpy(node->jacobian->x, A->x, A->nnz * sizeof(double));
+
+        /* Signal fast path by leaving CSC_work as NULL */
+        lin_node->CSC_work = NULL;
+        return;
+    }
+
+    /* Slow path: general case - compute A @ J_x via matrix multiplication */
+    lin_node->CSC_work = csr_to_csc_fill_sparsity(x->jacobian, node->iwork);
     node->jacobian = csr_csc_matmul_alloc(lin_node->A, lin_node->CSC_work);
 }
 
@@ -77,11 +103,18 @@ static void eval_jacobian(expr *node)
     expr *x = node->left;
     left_matmul_expr *lin_node = (left_matmul_expr *) node;
 
-    /* evaluate child's jacobian and convert to CSC */
+    /* evaluate child's jacobian */
     x->eval_jacobian(x);
-    csr_to_csc_fill_values(x->jacobian, lin_node->CSC_work, node->iwork);
 
-    /* compute this node's jacobian */
+    /* Fast path: child is a variable.
+     * Values already set in jacobian_init (A is constant, variable Jacobian is 1.0) */
+    if (lin_node->CSC_work == NULL)
+    {
+        return;
+    }
+
+    /* Slow path: convert to CSC and multiply */
+    csr_to_csc_fill_values(x->jacobian, lin_node->CSC_work, node->iwork);
     csr_csc_matmul_fill_values(lin_node->A, lin_node->CSC_work, node->jacobian);
 }
 

src/bivariate/right_matmul.c

@@ -54,11 +54,38 @@ static void jacobian_init(expr *node)
     expr *x = node->left;
     right_matmul_expr *right_node = (right_matmul_expr *) node;
 
-    /* initialize child's jacobian and precompute sparsity of its transpose */
+    /* initialize child's jacobian */
     x->jacobian_init(x);
-    right_node->CSC_work = csr_to_csc_fill_sparsity(x->jacobian, node->iwork);
 
-    /* precompute sparsity of this node's jacobian */
+    /* Fast path: child is a variable with identity Jacobian.
+     * For y = x @ A, we compute vec(y) = B @ vec(x) where B = A^T kron I_p.
+     * J_y = B @ J_x where J_x is identity shifted by var_id.
+     * So J_y is just B with column indices shifted by var_id.
+     * This avoids O(nnz(B) * n) CSR-CSC multiplication. */
+    if (x->var_id != NOT_A_VARIABLE)
+    {
+        CSR_Matrix *B = right_node->B;
+        node->jacobian = new_csr_matrix(B->m, node->n_vars, B->nnz);
+
+        /* Copy row pointers directly */
+        memcpy(node->jacobian->p, B->p, (B->m + 1) * sizeof(int));
+
+        /* Copy column indices with offset by var_id */
+        for (int k = 0; k < B->nnz; k++)
+        {
+            node->jacobian->i[k] = B->i[k] + x->var_id;
+        }
+
+        /* Copy values directly (B is constant, variable Jacobian entries are 1.0) */
+        memcpy(node->jacobian->x, B->x, B->nnz * sizeof(double));
+
+        /* Signal fast path by leaving CSC_work as NULL */
+        right_node->CSC_work = NULL;
+        return;
+    }
+
+    /* Slow path: general case - compute B @ J_x via matrix multiplication */
+    right_node->CSC_work = csr_to_csc_fill_sparsity(x->jacobian, node->iwork);
     node->jacobian = csr_csc_matmul_alloc(right_node->B, right_node->CSC_work);
 }
 
@@ -67,11 +94,18 @@ static void eval_jacobian(expr *node)
     expr *x = node->left;
     right_matmul_expr *right_node = (right_matmul_expr *) node;
 
-    /* evaluate child's jacobian and convert to CSC*/
+    /* evaluate child's jacobian */
     x->eval_jacobian(x);
-    csr_to_csc_fill_values(x->jacobian, right_node->CSC_work, node->iwork);
 
-    /* compute this node's jacobian */
+    /* Fast path: child is a variable.
+     * Values already set in jacobian_init (B is constant, variable Jacobian is 1.0) */
+    if (right_node->CSC_work == NULL)
+    {
+        return;
+    }
+
+    /* Slow path: convert to CSC and multiply */
+    csr_to_csc_fill_values(x->jacobian, right_node->CSC_work, node->iwork);
     csr_csc_matmul_fill_values(right_node->B, right_node->CSC_work, node->jacobian);
 }
 

src/dnlp_diff_engine/__init__.py

@@ -68,8 +68,8 @@ def _convert_multiply(expr, children):
     # multiply has args: [left, right]
     left_arg, right_arg = expr.args
 
-    # Check if left is a constant
-    if isinstance(left_arg, cp.Constant):
+    # Check if left is a constant (use is_constant() to handle Promote and other wrappers)
+    if left_arg.is_constant():
         value = np.asarray(left_arg.value, dtype=np.float64)
 
         # Scalar constant
@@ -83,7 +83,7 @@ def _convert_multiply(expr, children):
             return _diffengine.make_const_vector_mult(children[1], vector)
 
     # Check if right is a constant
-    elif isinstance(right_arg, cp.Constant):
+    elif right_arg.is_constant():
         value = np.asarray(right_arg.value, dtype=np.float64)
 
         # Scalar constant
@@ -323,10 +323,15 @@ def __init__(self, cvxpy_problem: cp.Problem):
         self._allocated = False
 
     def init_derivatives(self):
-        """Initialize derivative structures. Must be called before forward/gradient/jacobian."""
+        """Initialize derivative structures. Must be called before forward/gradient/jacobian/hessian."""
         _diffengine.problem_init_derivatives(self._capsule)
         self._allocated = True
 
+    def init_jacobian_only(self):
+        """Initialize Jacobian structure only (skip Hessian). Use when only Jacobian is needed."""
+        _diffengine.problem_init_jacobian_only(self._capsule)
+        self._allocated = True
+
     def objective_forward(self, u: np.ndarray) -> float:
         """Evaluate objective. Returns obj_value float."""
         return _diffengine.problem_objective_forward(self._capsule, u)

src/problem.c

@@ -104,6 +104,27 @@ void problem_init_derivatives(problem *prob)
     free(iwork);
 }
 
+void problem_init_jacobian_only(problem *prob)
+{
+    // -------------------------------------------------------------------------------
+    //                           Jacobian structure only (skip Hessian)
+    // -------------------------------------------------------------------------------
+    prob->objective->jacobian_init(prob->objective);
+    int nnz = 0;
+    for (int i = 0; i < prob->n_constraints; i++)
+    {
+        expr *c = prob->constraints[i];
+        c->jacobian_init(c);
+        nnz += c->jacobian->nnz;
+    }
+
+    prob->jacobian = new_csr_matrix(prob->total_constraint_size, prob->n_vars, nnz);
+
+    /* Leave Hessian structures as NULL */
+    prob->lagrange_hessian = NULL;
+    prob->hess_idx_map = NULL;
+}
+
 static void problem_lagrange_hess_fill_sparsity(problem *prob, int *iwork)
 {
     expr **constrs = prob->constraints;