fix(differentiable): support vector / mixed-shape constraints (closes #324)

[jkitchin]

Summary

Fixes #324. differentiable_solve's envelope-theorem sensitivity step assembled per-constraint values with jnp.array([cf(...) for cf in constraint_fns]) — a stack, valid only when every constraint is scalar. The multipliers vector from the NLP solve is per constraint row (the evaluator solves over concat_constraints, which ravels each constraint with reshape(-1)), so the stack mismatches.

Broader than the issue title: the path was broken for any non-scalar constraint, not just mixed shapes — mixed shapes fail to concatenate (the reported TypeError), and even equal shapes mis-align the λ·g dot (mults length 6 vs con_vals shape (2,3)).

Fix

Ravel + concatenate each constraint value, matching the per-row multiplier ordering:

con_vals = jnp.concatenate([jnp.ravel(cf(...)) for cf in constraint_fns])

Verified the ordering is correct (not just non-crashing): differentiable.py and NLPEvaluator both build constraint_fns via [c for c in model._constraints if isinstance(c, Constraint)] (same filter/order), and the evaluator ravels with reshape(-1) — so multipliers aligns row-for-row. For scalar constraints the new code is identical to the old.

Five sites had the pattern, all fixed: differentiable_solve, _compute_sensitivity_at_solution, solve_fn_jvp (custom-JVP), implicit_differentiate, and differentiable_solve_l3. The issue named only the L1 path, but the L3 path crashed identically — a regression test caught it (my first pass missed two sites; the test exposed them).

Tests

test_differentiable_vector_constraints.py (6 tests):

• mixed-shape (3,)+(2,) and equal-shape (3,)+(3,) no longer crash, on both the L1 and L3 paths;
• the gradient is checked correct against a central finite difference on a binding model (analytic d(obj)/dp = −2.1, exact match) — guarding both the crash and a silent multiplier-ordering regression (which would return a wrong gradient, not crash);
• scalar control unregressed.

Full differentiable suite (78) passes; ruff/mypy clean.

🤖 Generated with Claude Code

https://claude.ai/code/session_01Lq9NRuyCAWVxWoFv6AFuAt