Split converters.py into helpers.py, registry.py, and converters.py

- helpers.py: shared utilities, matmul helpers, var/param dict builders
- registry.py: atom converter functions and ATOM_CONVERTERS dict
- converters.py: convert_expr entry point and param-aware matmul/multiply

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Author: Transurgeon

cvxpy/reductions/solvers/nlp_solvers/diff_engine/__init__.py

@@ -22,12 +22,12 @@
 """
 
 from cvxpy.reductions.solvers.nlp_solvers.diff_engine.c_problem import C_problem
-from cvxpy.reductions.solvers.nlp_solvers.diff_engine.converters import (
-    ATOM_CONVERTERS,
+from cvxpy.reductions.solvers.nlp_solvers.diff_engine.converters import convert_expr
+from cvxpy.reductions.solvers.nlp_solvers.diff_engine.helpers import (
     build_param_dict,
     build_var_dict,
-    convert_expr,
 )
+from cvxpy.reductions.solvers.nlp_solvers.diff_engine.registry import ATOM_CONVERTERS
 
 __all__ = [
     "C_problem",

cvxpy/reductions/solvers/nlp_solvers/diff_engine/c_problem.py

@@ -17,10 +17,10 @@
 
 import cvxpy as cp
 from cvxpy.reductions.inverse_data import InverseData
-from cvxpy.reductions.solvers.nlp_solvers.diff_engine.converters import (
+from cvxpy.reductions.solvers.nlp_solvers.diff_engine.converters import convert_expr
+from cvxpy.reductions.solvers.nlp_solvers.diff_engine.helpers import (
     build_param_dict,
     build_var_dict,
-    convert_expr,
 )
 
 

cvxpy/reductions/solvers/nlp_solvers/diff_engine/converters.py

@@ -13,105 +13,25 @@
 See the License for the specific language governing permissions and
 limitations under the License.
 
-Converters from CVXPY expressions to C diff engine expressions.
-
-This module provides the mapping between CVXPY atom types and their
-corresponding SparseDiffPy constructors.
+Main entry point for converting CVXPY expressions to C diff engine expressions.
 """
-import numpy as np
 from scipy import sparse
 from sparsediffpy import _sparsediffengine as _diffengine
 
 import cvxpy as cp
 from cvxpy.expressions.constants.parameter import Parameter
 
-# ---------------------------------------------------------------------------
-# Utilities
-# ---------------------------------------------------------------------------
-
-def normalize_shape(shape):
-    """Normalize shape to 2D (d1, d2) for the C engine."""
-    shape = tuple(shape)
-    return (1,) * (2 - len(shape)) + shape
-
-
-def _to_dense_float(value):
-    """Convert a value to a dense float64 numpy array."""
-    if sparse.issparse(value):
-        value = value.todense()
-    return np.asarray(value, dtype=np.float64)
-
-
-def _chain_add(children):
-    """Chain multiple children with binary adds: a + b + c -> add(add(a, b), c)."""
-    result = children[0]
-    for child in children[1:]:
-        result = _diffengine.make_add(result, child)
-    return result
-
-
-# ---------------------------------------------------------------------------
-# Matmul helpers (handle param_node insertion for backward compat)
-# ---------------------------------------------------------------------------
-
-def _make_sparse_left_matmul(param_node, child, A):
-    if not isinstance(A, sparse.csr_matrix):
-        A = sparse.csr_matrix(A)
-    return _diffengine.make_sparse_left_matmul(
-        param_node, child,
-        A.data.astype(np.float64, copy=False),
-        A.indices.astype(np.int32, copy=False),
-        A.indptr.astype(np.int32, copy=False),
-        A.shape[0], A.shape[1])
-
-
-def _make_dense_left_matmul(param_node, child, A):
-    m, n = normalize_shape(A.shape)
-    return _diffengine.make_dense_left_matmul(
-        param_node, child, A.flatten(order='C'), m, n)
-
-
-def _make_sparse_right_matmul(param_node, child, A):
-    if not isinstance(A, sparse.csr_matrix):
-        A = sparse.csr_matrix(A)
-    return _diffengine.make_sparse_right_matmul(
-        param_node, child,
-        A.data.astype(np.float64, copy=False),
-        A.indices.astype(np.int32, copy=False),
-        A.indptr.astype(np.int32, copy=False),
-        A.shape[0], A.shape[1])
-
-
-def _make_dense_right_matmul(param_node, child, A):
-    m, n = normalize_shape(A.shape)
-    return _diffengine.make_dense_right_matmul(
-        param_node, child, A.flatten(order='C'), m, n)
-
-
-# ---------------------------------------------------------------------------
-# Variable / parameter dict builders
-# ---------------------------------------------------------------------------
-
-def build_var_dict(inverse_data):
-    """Build {var_id: C variable capsule} mapping from InverseData."""
-    n_vars = inverse_data.x_length
-    var_dict = {}
-    for var_id, (offset, _) in inverse_data.id_map.items():
-        d1, d2 = normalize_shape(inverse_data.var_shapes[var_id])
-        var_dict[var_id] = _diffengine.make_variable(d1, d2, offset, n_vars)
-    return var_dict, n_vars
-
-
-def build_param_dict(inverse_data):
-    """Build {param_id: C parameter capsule} mapping from InverseData."""
-    n_vars = inverse_data.x_length
-    param_dict = {}
-    for param_id, offset in inverse_data.param_id_map.items():
-        if param_id not in inverse_data.param_shapes:
-            continue
-        d1, d2 = normalize_shape(inverse_data.param_shapes[param_id])
-        param_dict[param_id] = _diffengine.make_parameter(d1, d2, offset, n_vars)
-    return param_dict
+from cvxpy.reductions.solvers.nlp_solvers.diff_engine.helpers import (
+    _make_dense_left_matmul,
+    _make_dense_right_matmul,
+    _make_sparse_left_matmul,
+    _make_sparse_right_matmul,
+    _to_dense_float,
+    build_param_dict,
+    build_var_dict,
+    normalize_shape,
+)
+from cvxpy.reductions.solvers.nlp_solvers.diff_engine.registry import ATOM_CONVERTERS
 
 
 # ---------------------------------------------------------------------------
@@ -186,218 +106,6 @@ def _convert_multiply(expr, children, var_dict, n_vars, param_dict):
     return _diffengine.make_multiply(children[0], children[1])
 
 
-# ---------------------------------------------------------------------------
-# Atom converters: (expr, children) -> C expression
-# These are kept exactly as in the original code.
-# ---------------------------------------------------------------------------
-
-def _convert_hstack(expr, children):
-    """Convert horizontal stack (hstack) of expressions."""
-    return _diffengine.make_hstack(children)
-
-
-def _extract_flat_indices_from_index(expr):
-    """Extract flattened indices from CVXPY index expression."""
-    parent_shape = expr.args[0].shape
-    indices_per_dim = [np.arange(s.start, s.stop, s.step) for s in expr.key]
-
-    if len(indices_per_dim) == 1:
-        return indices_per_dim[0].astype(np.int32)
-    elif len(indices_per_dim) == 2:
-        # Fortran order: idx = row + col * n_rows
-        return (
-            np.add.outer(indices_per_dim[0], indices_per_dim[1] * parent_shape[0])
-            .flatten(order="F")
-            .astype(np.int32)
-        )
-    else:
-        raise NotImplementedError("index with >2 dimensions not supported")
-
-
-def _extract_flat_indices_from_special_index(expr):
-    """Extract flattened indices from CVXPY special_index expression."""
-    return np.reshape(expr._select_mat, expr._select_mat.size, order="F").astype(np.int32)
-
-
-def _convert_rel_entr(expr, children):
-    """Convert rel_entr(x, y) = x * log(x/y) elementwise.
-
-    Uses specialized functions based on argument shapes:
-    - Both scalar or both same size: make_rel_entr (elementwise)
-    - First arg vector, second scalar: make_rel_entr_vector_scalar
-    - First arg scalar, second vector: make_rel_entr_scalar_vector
-    """
-    x_arg, y_arg = expr.args
-    x_size = x_arg.size
-    y_size = y_arg.size
-
-    # Determine which variant to use based on sizes
-    if x_size == y_size:
-        return _diffengine.make_rel_entr(children[0], children[1])
-    elif x_size > 1 and y_size == 1:
-        return _diffengine.make_rel_entr_vector_scalar(children[0], children[1])
-    elif x_size == 1 and y_size > 1:
-        return _diffengine.make_rel_entr_scalar_vector(children[0], children[1])
-    else:
-        raise ValueError(
-            f"rel_entr requires arguments to be either both scalars, both same size, "
-            f"or one scalar and one vector. Got sizes: x={x_size}, y={y_size}"
-        )
-
-
-def _convert_quad_form(expr, children):
-    """Convert quadratic form x.T @ P @ x."""
-
-    P = expr.args[1]
-
-    if not isinstance(P, cp.Constant):
-        raise NotImplementedError("quad_form requires P to be a constant matrix")
-
-    P = P.value
-
-    if not isinstance(P, sparse.csr_matrix):
-          P = sparse.csr_matrix(P)
-
-    return _diffengine.make_quad_form(
-        children[0],
-        P.data.astype(np.float64),
-        P.indices.astype(np.int32),
-        P.indptr.astype(np.int32),
-        P.shape[0],
-        P.shape[1],
-    )
-
-
-def _convert_reshape(expr, children):
-    """Convert reshape - only Fortran order is supported.
-
-    Note: Only order='F' (Fortran/column-major) is supported.
-    """
-    if expr.order != "F":
-        raise NotImplementedError(
-            f"reshape with order='{expr.order}' not supported. "
-            "Only order='F' (Fortran) is currently supported."
-        )
-
-    d1, d2 = normalize_shape(expr.shape)
-    return _diffengine.make_reshape(children[0], d1, d2)
-
-def _convert_broadcast(expr, children):
-    d1, d2 = expr.broadcast_shape
-    d1_C, d2_C = _diffengine.get_expr_dimensions(children[0])
-    if d1_C == d1 and d2_C == d2:
-        return children[0]
-
-    return _diffengine.make_broadcast(children[0], d1, d2)
-
-def _convert_sum(expr, children):
-    axis = expr.axis
-    if axis is None:
-        axis = -1
-    return _diffengine.make_sum(children[0], axis)
-
-def _convert_promote(expr, children):
-    d1, d2 = normalize_shape(expr.shape)
-    return _diffengine.make_promote(children[0], d1, d2)
-
-def _convert_NegExpression(_expr, children):
-    return _diffengine.make_neg(children[0])
-
-def _convert_quad_over_lin(_expr, children):
-    return _diffengine.make_quad_over_lin(children[0], children[1])
-
-def _convert_index(expr, children):
-    idxs = _extract_flat_indices_from_index(expr)
-    d1, d2 = normalize_shape(expr.shape)
-    return _diffengine.make_index(children[0], d1, d2, idxs)
-
-def _convert_special_index(expr, children):
-    idxs = _extract_flat_indices_from_special_index(expr)
-    d1, d2 = normalize_shape(expr.shape)
-    return _diffengine.make_index(children[0], d1, d2, idxs)
-
-def _convert_prod(expr, children):
-    axis = expr.axis
-    if axis is None:
-        return _diffengine.make_prod(children[0])
-    elif axis == 0:
-        return _diffengine.make_prod_axis_zero(children[0])
-    elif axis == 1:
-        return _diffengine.make_prod_axis_one(children[0])
-
-def _convert_transpose(expr, children):
-    # If the child is a vector (shape (n,) or (n,1) or (1,n)), use reshape to transpose
-    child_shape = normalize_shape(expr.args[0].shape)
-
-    if 1 in child_shape:
-        return _diffengine.make_reshape(children[0], child_shape[1], child_shape[0])
-    else:
-        return _diffengine.make_transpose(children[0])
-
-def _convert_trace(_expr, children):
-    return _diffengine.make_trace(children[0])
-
-def _convert_diag_vec(expr, children):
-    # C implementation only supports k=0 (main diagonal)
-    if expr.k != 0:
-        raise NotImplementedError("diag_vec with k != 0 not supported in diff engine")
-    return _diffengine.make_diag_vec(children[0])
-
-
-# ---------------------------------------------------------------------------
-# Atom converter registry
-# Converters receive (expr, children) where expr is the CVXPY expression.
-# matmul and multiply are handled separately (they need param_dict).
-# ---------------------------------------------------------------------------
-
-ATOM_CONVERTERS = {
-    # Elementwise unary
-    "log": lambda _expr, children: _diffengine.make_log(children[0]),
-    "exp": lambda _expr, children: _diffengine.make_exp(children[0]),
-    # Affine unary
-    "NegExpression": _convert_NegExpression,
-    "Promote": _convert_promote,
-    # N-ary (handles 2+ args)
-    "AddExpression": lambda _expr, children: _chain_add(children),
-    # Reductions
-    "Sum": _convert_sum,
-    # Bivariate
-    "QuadForm": _convert_quad_form,
-    "quad_over_lin": _convert_quad_over_lin,
-    "rel_entr": _convert_rel_entr,
-    # Elementwise univariate with parameter
-    "Power": lambda expr, children: _diffengine.make_power(children[0], float(expr.p.value)),
-    "PowerApprox": lambda expr, children: _diffengine.make_power(children[0], float(expr.p.value)),
-    # Trigonometric
-    "sin": lambda _expr, children: _diffengine.make_sin(children[0]),
-    "cos": lambda _expr, children: _diffengine.make_cos(children[0]),
-    "tan": lambda _expr, children: _diffengine.make_tan(children[0]),
-    # Hyperbolic
-    "sinh": lambda _expr, children: _diffengine.make_sinh(children[0]),
-    "tanh": lambda _expr, children: _diffengine.make_tanh(children[0]),
-    "asinh": lambda _expr, children: _diffengine.make_asinh(children[0]),
-    "atanh": lambda _expr, children: _diffengine.make_atanh(children[0]),
-    # Other elementwise
-    "entr": lambda _expr, children: _diffengine.make_entr(children[0]),
-    "logistic": lambda _expr, children: _diffengine.make_logistic(children[0]),
-    "xexp": lambda _expr, children: _diffengine.make_xexp(children[0]),
-    "normcdf": lambda _expr, children: _diffengine.make_normal_cdf(children[0]),
-    # Indexing/slicing
-    "index": _convert_index,
-    "special_index": _convert_special_index,
-    "reshape": _convert_reshape,
-    "broadcast_to": _convert_broadcast,
-    # Reductions returning scalar
-    "Prod": _convert_prod,
-    "transpose": _convert_transpose,
-    # Horizontal stack
-    "Hstack": _convert_hstack,
-    "Trace": _convert_trace,
-    # Diagonal
-    "diag_vec": _convert_diag_vec,
-}
-
-
 # ---------------------------------------------------------------------------
 # Main conversion entry point
 # ---------------------------------------------------------------------------