elastiqp 0.2.0__tar.gz → 0.2.1__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (52) hide show
  1. {elastiqp-0.2.0 → elastiqp-0.2.1}/.github/workflows/wheels.yml +1 -1
  2. {elastiqp-0.2.0 → elastiqp-0.2.1}/PKG-INFO +1 -1
  3. {elastiqp-0.2.0 → elastiqp-0.2.1}/bindings/bindings.cc +1 -0
  4. {elastiqp-0.2.0 → elastiqp-0.2.1}/bindings/jax_ffi.cc +25 -68
  5. {elastiqp-0.2.0 → elastiqp-0.2.1}/include/elastiqp/das.hpp +36 -8
  6. {elastiqp-0.2.0 → elastiqp-0.2.1}/include/elastiqp/elastiqp_c.h +4 -1
  7. {elastiqp-0.2.0 → elastiqp-0.2.1}/package.xml +1 -1
  8. {elastiqp-0.2.0 → elastiqp-0.2.1}/python/elastiqp/jax.py +47 -135
  9. {elastiqp-0.2.0 → elastiqp-0.2.1}/src/elastiqp_c.cc +6 -2
  10. {elastiqp-0.2.0 → elastiqp-0.2.1}/tests/test_c_api.c +1 -1
  11. {elastiqp-0.2.0 → elastiqp-0.2.1}/tests/test_c_parity.cc +1 -0
  12. {elastiqp-0.2.0 → elastiqp-0.2.1}/tests/test_das.cc +57 -3
  13. {elastiqp-0.2.0 → elastiqp-0.2.1}/tests/test_jax_ffi.py +99 -29
  14. {elastiqp-0.2.0 → elastiqp-0.2.1}/.clang-format +0 -0
  15. {elastiqp-0.2.0 → elastiqp-0.2.1}/.gitignore +0 -0
  16. {elastiqp-0.2.0 → elastiqp-0.2.1}/CMakeLists.txt +0 -0
  17. {elastiqp-0.2.0 → elastiqp-0.2.1}/LICENSE +0 -0
  18. {elastiqp-0.2.0 → elastiqp-0.2.1}/NOTICE +0 -0
  19. {elastiqp-0.2.0 → elastiqp-0.2.1}/README.md +0 -0
  20. {elastiqp-0.2.0 → elastiqp-0.2.1}/cmake/elastiqpConfig.cmake.in +0 -0
  21. {elastiqp-0.2.0 → elastiqp-0.2.1}/cmake/elastiqp_install.cmake +0 -0
  22. {elastiqp-0.2.0 → elastiqp-0.2.1}/examples/basic_usage.py +0 -0
  23. {elastiqp-0.2.0 → elastiqp-0.2.1}/examples/constraint_conflict_demo.py +0 -0
  24. {elastiqp-0.2.0 → elastiqp-0.2.1}/examples/jax_usage.py +0 -0
  25. {elastiqp-0.2.0 → elastiqp-0.2.1}/examples/torch_usage.py +0 -0
  26. {elastiqp-0.2.0 → elastiqp-0.2.1}/include/elastiqp/common.hpp +0 -0
  27. {elastiqp-0.2.0 → elastiqp-0.2.1}/include/elastiqp/elastiqp.hpp +0 -0
  28. {elastiqp-0.2.0 → elastiqp-0.2.1}/include/elastiqp/ipm.hpp +0 -0
  29. {elastiqp-0.2.0 → elastiqp-0.2.1}/include/elastiqp/kkt_vjp.hpp +0 -0
  30. {elastiqp-0.2.0 → elastiqp-0.2.1}/include/elastiqp/pdal.hpp +0 -0
  31. {elastiqp-0.2.0 → elastiqp-0.2.1}/pyproject.toml +0 -0
  32. {elastiqp-0.2.0 → elastiqp-0.2.1}/python/elastiqp/__init__.py +0 -0
  33. {elastiqp-0.2.0 → elastiqp-0.2.1}/python/elastiqp/torch.py +0 -0
  34. {elastiqp-0.2.0 → elastiqp-0.2.1}/tests/data/warmstart_dual_overshoot.npz +0 -0
  35. {elastiqp-0.2.0 → elastiqp-0.2.1}/tests/data/warmstart_stale_duals.npz +0 -0
  36. {elastiqp-0.2.0 → elastiqp-0.2.1}/tests/support/drift_traj.hpp +0 -0
  37. {elastiqp-0.2.0 → elastiqp-0.2.1}/tests/support/problem_gen.hpp +0 -0
  38. {elastiqp-0.2.0 → elastiqp-0.2.1}/tests/support/test_util.hpp +0 -0
  39. {elastiqp-0.2.0 → elastiqp-0.2.1}/tests/test_bcl_creep.cc +0 -0
  40. {elastiqp-0.2.0 → elastiqp-0.2.1}/tests/test_bindings.py +0 -0
  41. {elastiqp-0.2.0 → elastiqp-0.2.1}/tests/test_das_alloc.cc +0 -0
  42. {elastiqp-0.2.0 → elastiqp-0.2.1}/tests/test_gap_creep.cc +0 -0
  43. {elastiqp-0.2.0 → elastiqp-0.2.1}/tests/test_ipm.cc +0 -0
  44. {elastiqp-0.2.0 → elastiqp-0.2.1}/tests/test_lp.cc +0 -0
  45. {elastiqp-0.2.0 → elastiqp-0.2.1}/tests/test_pdal.cc +0 -0
  46. {elastiqp-0.2.0 → elastiqp-0.2.1}/tests/test_ruiz.cc +0 -0
  47. {elastiqp-0.2.0 → elastiqp-0.2.1}/tests/test_solvers.cc +0 -0
  48. {elastiqp-0.2.0 → elastiqp-0.2.1}/tests/test_torch.py +0 -0
  49. {elastiqp-0.2.0 → elastiqp-0.2.1}/third_party/xla_ffi/README.md +0 -0
  50. {elastiqp-0.2.0 → elastiqp-0.2.1}/third_party/xla_ffi/xla/ffi/api/api.h +0 -0
  51. {elastiqp-0.2.0 → elastiqp-0.2.1}/third_party/xla_ffi/xla/ffi/api/c_api.h +0 -0
  52. {elastiqp-0.2.0 → elastiqp-0.2.1}/third_party/xla_ffi/xla/ffi/api/ffi.h +0 -0
@@ -31,7 +31,7 @@ jobs:
31
31
  strategy:
32
32
  fail-fast: false
33
33
  matrix:
34
- os: [ubuntu-latest, ubuntu-24.04-arm, macos-14, windows-latest]
34
+ os: [ubuntu-latest, ubuntu-24.04-arm, macos-15, windows-latest]
35
35
  steps:
36
36
  - uses: actions/checkout@v5
37
37
 
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.1
2
2
  Name: elastiqp
3
- Version: 0.2.0
3
+ Version: 0.2.1
4
4
  Summary: An always-feasible QP solver for constrained robot control
5
5
  Author-Email: Daniel Morton <danielpmorton@gmail.com>
6
6
  License: Copyright 2026 Daniel Morton
@@ -285,6 +285,7 @@ NB_MODULE(_core, m) {
285
285
  .def_rw("eps_rel", &S::eps_rel)
286
286
  .def_rw("sing_tol", &S::sing_tol)
287
287
  .def_rw("zero_tol", &S::zero_tol)
288
+ .def_rw("cond_iters", &S::cond_iters)
288
289
  .def_rw("eps_prox", &S::eps_prox)
289
290
  .def_rw("eta_prox", &S::eta_prox)
290
291
  .def_rw("prox_relaxation", &S::prox_relaxation)
@@ -1,16 +1,16 @@
1
1
  // XLA FFI handlers for ElastiQP. Every call builds a fresh solver (XLA
2
2
  // assumes purity), so warm starting is an explicit argument.
3
3
  //
4
- // ElastiqpSolve: (Q, q, A, b, G, h, penalty)
5
- // -> (x, t, y, z_t, z, xr, tr, yr, z_t_r, z_r, info)
6
- // ElastiqpSolveWarm: (Q, q, A, b, G, h, penalty, x0, y0, z0)
7
- // -> (x, t, y, z_t, z, info) (das and pdal only)
4
+ // ElastiqpSolve: (Q, q, A, b, G, h, penalty, x0, y0, z0)
5
+ // -> (x, t, y, z_t, z, xr, tr, yr, z_t_r, z_r, info)
8
6
  //
9
7
  // method: 0 = das, 1 = pdal, 2 = ipm. max_iter is the backend's outer budget.
8
+ // warm != 0 seeds the solve from (x0, y0, z0) (das and pdal only); otherwise
9
+ // they are ignored.
10
10
  // The *r block is the kappa-relaxed point the Python wrappers differentiate
11
11
  // at; it copies the tight solution when target_kappa <= 0. info = [converged,
12
- // iters, relax_converged], without the last entry for the warm handler. A has 0
13
- // rows and b 0 entries when there are no equalities.
12
+ // iters, relax_converged]. A has 0 rows and b 0 entries when there are no
13
+ // equalities.
14
14
 
15
15
  #include <algorithm>
16
16
  #include <cstdint>
@@ -134,9 +134,11 @@ void write_solution(const elastiqp::Solution& sol, int64_t n, int64_t m,
134
134
 
135
135
  ffi::Error ElastiqpSolveImpl(
136
136
  double eps_abs, int64_t max_iter, int64_t ruiz, double target_kappa,
137
- int64_t method, ffi::Buffer<ffi::F64> Q, ffi::Buffer<ffi::F64> q,
138
- ffi::Buffer<ffi::F64> A, ffi::Buffer<ffi::F64> b, ffi::Buffer<ffi::F64> G,
139
- ffi::Buffer<ffi::F64> h, ffi::Buffer<ffi::F64> penalty,
137
+ int64_t warm, int64_t method, ffi::Buffer<ffi::F64> Q,
138
+ ffi::Buffer<ffi::F64> q, ffi::Buffer<ffi::F64> A, ffi::Buffer<ffi::F64> b,
139
+ ffi::Buffer<ffi::F64> G, ffi::Buffer<ffi::F64> h,
140
+ ffi::Buffer<ffi::F64> penalty, ffi::Buffer<ffi::F64> x0,
141
+ ffi::Buffer<ffi::F64> y0, ffi::Buffer<ffi::F64> z0,
140
142
  ffi::ResultBuffer<ffi::F64> x, ffi::ResultBuffer<ffi::F64> t,
141
143
  ffi::ResultBuffer<ffi::F64> y, ffi::ResultBuffer<ffi::F64> z_t,
142
144
  ffi::ResultBuffer<ffi::F64> z, ffi::ResultBuffer<ffi::F64> xr,
@@ -146,37 +148,9 @@ ffi::Error ElastiqpSolveImpl(
146
148
  const int64_t n = q.dimensions()[0];
147
149
  const int64_t m = b.dimensions()[0];
148
150
  const int64_t p = h.dimensions()[0];
149
-
150
- elastiqp::Solution sol, rsol;
151
- ffi::Error err =
152
- dispatch(method, eps_abs, max_iter, ruiz, target_kappa, n, m, p,
153
- Q.typed_data(), q.typed_data(), A.typed_data(), b.typed_data(),
154
- G.typed_data(), h.typed_data(), penalty.typed_data(), nullptr,
155
- nullptr, nullptr, sol, rsol);
156
- if (err.failure()) return err;
157
-
158
- write_solution(sol, n, m, p, x, t, y, z_t, z);
159
- write_solution(rsol, n, m, p, xr, tr, yr, z_t_r, z_r);
160
- info->typed_data()[0] = static_cast<double>(sol.converged);
161
- info->typed_data()[1] = static_cast<double>(sol.iters);
162
- info->typed_data()[2] = static_cast<double>(rsol.converged);
163
- return ffi::Error::Success();
164
- }
165
-
166
- ffi::Error ElastiqpSolveWarmImpl(
167
- double eps_abs, int64_t max_iter, int64_t ruiz, int64_t method,
168
- ffi::Buffer<ffi::F64> Q, ffi::Buffer<ffi::F64> q, ffi::Buffer<ffi::F64> A,
169
- ffi::Buffer<ffi::F64> b, ffi::Buffer<ffi::F64> G, ffi::Buffer<ffi::F64> h,
170
- ffi::Buffer<ffi::F64> penalty, ffi::Buffer<ffi::F64> x0,
171
- ffi::Buffer<ffi::F64> y0, ffi::Buffer<ffi::F64> z0,
172
- ffi::ResultBuffer<ffi::F64> x, ffi::ResultBuffer<ffi::F64> t,
173
- ffi::ResultBuffer<ffi::F64> y, ffi::ResultBuffer<ffi::F64> z_t,
174
- ffi::ResultBuffer<ffi::F64> z, ffi::ResultBuffer<ffi::F64> info) {
175
- const int64_t n = q.dimensions()[0];
176
- const int64_t m = b.dimensions()[0];
177
- const int64_t p = h.dimensions()[0];
178
- if (x0.dimensions()[0] != n || y0.dimensions()[0] != m ||
179
- z0.dimensions()[0] != p) {
151
+ const bool w = warm != 0;
152
+ if (w && (x0.dimensions()[0] != n || y0.dimensions()[0] != m ||
153
+ z0.dimensions()[0] != p)) {
180
154
  return ffi::Error(ffi::ErrorCode::kInvalidArgument,
181
155
  "warm start (x0, y0, z0) must have shapes (n,), (m,), "
182
156
  "(p,) matching (q, b, h)");
@@ -184,15 +158,18 @@ ffi::Error ElastiqpSolveWarmImpl(
184
158
 
185
159
  elastiqp::Solution sol, rsol;
186
160
  ffi::Error err =
187
- dispatch(method, eps_abs, max_iter, ruiz, 0.0, n, m, p, Q.typed_data(),
188
- q.typed_data(), A.typed_data(), b.typed_data(), G.typed_data(),
189
- h.typed_data(), penalty.typed_data(), x0.typed_data(),
190
- y0.typed_data(), z0.typed_data(), sol, rsol);
161
+ dispatch(method, eps_abs, max_iter, ruiz, target_kappa, n, m, p,
162
+ Q.typed_data(), q.typed_data(), A.typed_data(), b.typed_data(),
163
+ G.typed_data(), h.typed_data(), penalty.typed_data(),
164
+ w ? x0.typed_data() : nullptr, w ? y0.typed_data() : nullptr,
165
+ w ? z0.typed_data() : nullptr, sol, rsol);
191
166
  if (err.failure()) return err;
192
167
 
193
168
  write_solution(sol, n, m, p, x, t, y, z_t, z);
169
+ write_solution(rsol, n, m, p, xr, tr, yr, z_t_r, z_r);
194
170
  info->typed_data()[0] = static_cast<double>(sol.converged);
195
171
  info->typed_data()[1] = static_cast<double>(sol.iters);
172
+ info->typed_data()[2] = static_cast<double>(rsol.converged);
196
173
  return ffi::Error::Success();
197
174
  }
198
175
 
@@ -204,6 +181,7 @@ XLA_FFI_DEFINE_HANDLER_SYMBOL(ElastiqpSolve, ElastiqpSolveImpl,
204
181
  .Attr<int64_t>("max_iter")
205
182
  .Attr<int64_t>("ruiz")
206
183
  .Attr<double>("target_kappa")
184
+ .Attr<int64_t>("warm")
207
185
  .Attr<int64_t>("method")
208
186
  .Arg<ffi::Buffer<ffi::F64>>() // Q
209
187
  .Arg<ffi::Buffer<ffi::F64>>() // q
@@ -212,6 +190,9 @@ XLA_FFI_DEFINE_HANDLER_SYMBOL(ElastiqpSolve, ElastiqpSolveImpl,
212
190
  .Arg<ffi::Buffer<ffi::F64>>() // G
213
191
  .Arg<ffi::Buffer<ffi::F64>>() // h
214
192
  .Arg<ffi::Buffer<ffi::F64>>() // penalty
193
+ .Arg<ffi::Buffer<ffi::F64>>() // x0
194
+ .Arg<ffi::Buffer<ffi::F64>>() // y0
195
+ .Arg<ffi::Buffer<ffi::F64>>() // z0
215
196
  .Ret<ffi::Buffer<ffi::F64>>() // x
216
197
  .Ret<ffi::Buffer<ffi::F64>>() // t
217
198
  .Ret<ffi::Buffer<ffi::F64>>() // y
@@ -224,27 +205,3 @@ XLA_FFI_DEFINE_HANDLER_SYMBOL(ElastiqpSolve, ElastiqpSolveImpl,
224
205
  .Ret<ffi::Buffer<ffi::F64>>() // z_r
225
206
  .Ret<ffi::Buffer<ffi::F64>>() // info
226
207
  );
227
-
228
- XLA_FFI_DEFINE_HANDLER_SYMBOL(ElastiqpSolveWarm, ElastiqpSolveWarmImpl,
229
- ffi::Ffi::Bind()
230
- .Attr<double>("eps_abs")
231
- .Attr<int64_t>("max_iter")
232
- .Attr<int64_t>("ruiz")
233
- .Attr<int64_t>("method")
234
- .Arg<ffi::Buffer<ffi::F64>>() // Q
235
- .Arg<ffi::Buffer<ffi::F64>>() // q
236
- .Arg<ffi::Buffer<ffi::F64>>() // A
237
- .Arg<ffi::Buffer<ffi::F64>>() // b
238
- .Arg<ffi::Buffer<ffi::F64>>() // G
239
- .Arg<ffi::Buffer<ffi::F64>>() // h
240
- .Arg<ffi::Buffer<ffi::F64>>() // penalty
241
- .Arg<ffi::Buffer<ffi::F64>>() // x0
242
- .Arg<ffi::Buffer<ffi::F64>>() // y0
243
- .Arg<ffi::Buffer<ffi::F64>>() // z0
244
- .Ret<ffi::Buffer<ffi::F64>>() // x
245
- .Ret<ffi::Buffer<ffi::F64>>() // t
246
- .Ret<ffi::Buffer<ffi::F64>>() // y
247
- .Ret<ffi::Buffer<ffi::F64>>() // z_t
248
- .Ret<ffi::Buffer<ffi::F64>>() // z
249
- .Ret<ffi::Buffer<ffi::F64>>() // info
250
- );
@@ -21,8 +21,12 @@ struct Settings {
21
21
  double eps_rel = 0.0;
22
22
  // Working-set LDL^T pivot below this = dependent row.
23
23
  double sing_tol = 3.7e-11;
24
- // Min/max Cholesky pivot ratio accepted for Q_s before adding a prox shift.
24
+ // Min/max eigenvalue ratio accepted for Q_s before adding a prox shift,
25
+ // judged from the Cholesky pivots (see cond_iters).
25
26
  double zero_tol = 1e-11;
27
+ // Inverse-iteration steps confirming the smallest eigenvalue of Q_s against
28
+ // zero_tol when the pivot ratio is below sqrt(zero_tol); 0 = pivots only.
29
+ int cond_iters = 3;
26
30
  // Initial prox shift (x max|diag Q_s|) when Q_s is not PD; 0 disables.
27
31
  double eps_prox = 1e-6;
28
32
  // Outer-loop tolerance on eps*|x - xc|_inf; 0 -> eps_abs.
@@ -109,6 +113,7 @@ class Solver {
109
113
  fx_.resize(n_);
110
114
  fr_.resize(mp_);
111
115
  xu_.resize(n_);
116
+ eig_work_.resize(n_);
112
117
  changed_.reserve(static_cast<size_t>(mp_));
113
118
  changed_cols_.resize(n_, mp_);
114
119
  warm_x_.setZero(n_);
@@ -305,16 +310,23 @@ class Solver {
305
310
  int inner = 0;
306
311
  st = ldp(inner);
307
312
  sol_.iters += inner;
308
- if (st == Status::kNumerics && eps_ > 0 &&
313
+ if (st == Status::kNumerics && settings.eps_prox > 0 &&
309
314
  escalations < settings.prox_escalations) {
310
315
  // Inner loop broke down: recenter at the current x, grow the prox shift
311
- // 100x, restart.
316
+ // 100x (or turn it on), restart.
312
317
  escalations++;
313
318
  prox_escalated_ = true;
314
319
  for (int i = 0; i < static_cast<int>(W_.size()); ++i)
315
320
  lam_full_[W_[static_cast<size_t>(i)]] = lam_[i];
316
- xc_ = llt_.matrixU().solve(u_ - v_);
317
- if (!factor(100.0 * eps_)) return finish(Status::kNumerics);
321
+ double eps_next = 100.0 * eps_;
322
+ if (eps_ > 0) {
323
+ xc_ = llt_.matrixU().solve(u_ - v_);
324
+ } else {
325
+ xc_ = x_;
326
+ eps_next = settings.eps_prox *
327
+ std::max(1.0, Qs_.diagonal().cwiseAbs().maxCoeff());
328
+ }
329
+ if (!factor(eps_next)) return finish(Status::kNumerics);
318
330
  for (int i = 0; i < mp_; ++i)
319
331
  tol_[i] = (settings.eps_abs + settings.eps_rel * std::abs(rhs_[i])) *
320
332
  scale_[i] * dr_[i];
@@ -653,6 +665,9 @@ class Solver {
653
665
  pmax = std::max(pmax, piv * piv);
654
666
  }
655
667
  ok = pmin > settings.zero_tol * pmax && pmin > 0;
668
+ if (ok && settings.cond_iters > 0 &&
669
+ pmin < std::sqrt(settings.zero_tol) * pmax)
670
+ ok = min_eig_estimate() > settings.zero_tol * (scale + eps_);
656
671
  }
657
672
  if (ok) break;
658
673
  if (settings.eps_prox <= 0) return false;
@@ -667,6 +682,18 @@ class Solver {
667
682
  rhs_dirty_ = true;
668
683
  return true;
669
684
  }
685
+ // Upper estimate of the smallest eigenvalue of the factored matrix
686
+ double min_eig_estimate() {
687
+ for (int i = 0; i < n_; ++i)
688
+ eig_work_[i] = 1.0 + static_cast<double>(i) / n_;
689
+ double est = std::numeric_limits<double>::infinity();
690
+ for (int it = 0; it < settings.cond_iters; ++it) {
691
+ eig_work_ /= eig_work_.norm();
692
+ llt_.solveInPlace(eig_work_);
693
+ est = 1.0 / eig_work_.norm();
694
+ }
695
+ return est;
696
+ }
670
697
  void normalize_row(int i) {
671
698
  const double nrm = Mt_.col(i).norm();
672
699
  scale_[i] = nrm > 1e-300 ? 1.0 / nrm : 1.0;
@@ -1322,9 +1349,10 @@ class Solver {
1322
1349
  double eps_ = 0.0;
1323
1350
  MatrixXd Mt_;
1324
1351
  VectorXd scale_, hi_, d_, v_, u_, uS_, mu_, x_, xc_, lam_full_;
1325
- // Scratch: Ruiz factors, a user-frame / unscaled n-vector, the columns
1326
- // update_rows re-solves, and the rows refactor_working_set re-adds.
1327
- VectorXd fx_, fr_, xu_;
1352
+ // Scratch: Ruiz factors, a user-frame / unscaled n-vector, the inverse
1353
+ // iteration vector, the columns update_rows re-solves, and the rows
1354
+ // refactor_working_set re-adds.
1355
+ VectorXd fx_, fr_, xu_, eig_work_;
1328
1356
  MatrixXd changed_cols_;
1329
1357
  std::vector<int> changed_, rows_;
1330
1358
  std::vector<RowState> state_;
@@ -74,7 +74,8 @@ typedef struct {
74
74
  * elastiqp_option_info() gives their ranges. Out-of-range values are rejected.
75
75
  *
76
76
  * Tolerances and iteration limits apply on the next solve. Scaling and
77
- * factorization options (ruiz*, zero_tol, eps_prox, reuse_factorization)
77
+ * factorization options (ruiz*, zero_tol, cond_iters, eps_prox,
78
+ * reuse_factorization)
78
79
  * apply on refactorization; updating Q forces it. Disabling ruiz retains
79
80
  * existing scaling. Explicit warm starts apply even with warm_start = 0.
80
81
  */
@@ -105,6 +106,8 @@ typedef struct {
105
106
 
106
107
  double relax_reg;
107
108
  int relax_factor_retries;
109
+
110
+ int cond_iters; /* appended to keep the struct layout of earlier fields */
108
111
  } ElastiQPSettings;
109
112
 
110
113
  /*
@@ -2,7 +2,7 @@
2
2
  <?xml-model href="http://download.ros.org/schema/package_format3.xsd" schematypens="http://www.w3.org/2001/XMLSchema"?>
3
3
  <package format="3">
4
4
  <name>elastiqp</name>
5
- <version>0.2.0</version>
5
+ <version>0.2.1</version>
6
6
  <description>ElastiQP: An elastic QP solver for robot control</description>
7
7
  <maintainer email="danielpmorton@gmail.com">Daniel Morton</maintainer>
8
8
  <license>Apache 2.0</license>
@@ -8,7 +8,7 @@ Supported:
8
8
  Notes:
9
9
  - vmap is sequential mode
10
10
  - Requires float64
11
- - Warm-starting is not differentiable
11
+ - Warm-starting is differentiable (the warm start itself gets a zero gradient)
12
12
  - Warm-starting requires explicit passing of the previous result
13
13
  - Recommended default kappa for differentiability: 1e-3
14
14
  """
@@ -66,11 +66,6 @@ _lib = ctypes.cdll.LoadLibrary(str(_find_library()))
66
66
  jax.ffi.register_ffi_target(
67
67
  "elastiqp_solve", jax.ffi.pycapsule(_lib.ElastiqpSolve), platform="cpu"
68
68
  )
69
- jax.ffi.register_ffi_target(
70
- "elastiqp_solve_warm",
71
- jax.ffi.pycapsule(_lib.ElastiqpSolveWarm),
72
- platform="cpu",
73
- )
74
69
 
75
70
 
76
71
  class Result(NamedTuple):
@@ -91,10 +86,14 @@ def _ffi_solve(
91
86
  G,
92
87
  h,
93
88
  penalty,
89
+ x0,
90
+ y0,
91
+ z0,
94
92
  eps_abs,
95
93
  max_iter,
96
94
  ruiz,
97
95
  target_kappa,
96
+ warm,
98
97
  method,
99
98
  ):
100
99
  n = Q.shape[-1]
@@ -105,60 +104,9 @@ def _ffi_solve(
105
104
  # (x, t, y, z_t, z) tight solution, then the kappa-relaxed central
106
105
  # point (identical to the tight block when target_kappa <= 0), then
107
106
  # info = [converged, iters, relax_converged].
108
- out_types = [
109
- vec(n),
110
- vec(p),
111
- vec(m),
112
- vec(p),
113
- vec(p),
114
- vec(n),
115
- vec(p),
116
- vec(m),
117
- vec(p),
118
- vec(p),
119
- vec(3),
120
- ]
107
+ sol = [vec(n), vec(p), vec(m), vec(p), vec(p)]
108
+ out_types = sol + sol + [vec(3)]
121
109
  call = jax.ffi.ffi_call("elastiqp_solve", out_types, vmap_method="sequential")
122
- return call(
123
- Q,
124
- q,
125
- A,
126
- b,
127
- G,
128
- h,
129
- penalty,
130
- eps_abs=np.float64(eps_abs),
131
- max_iter=np.int64(max_iter),
132
- ruiz=np.int64(bool(ruiz)),
133
- target_kappa=np.float64(target_kappa),
134
- method=np.int64(_METHOD_ID[method]),
135
- )
136
-
137
-
138
- def _ffi_solve_warm(
139
- Q,
140
- q,
141
- A,
142
- b,
143
- G,
144
- h,
145
- penalty,
146
- x0,
147
- y0,
148
- z0,
149
- eps_abs,
150
- max_iter,
151
- ruiz,
152
- method,
153
- ):
154
- n = Q.shape[-1]
155
- m = b.shape[-1]
156
- p = h.shape[-1]
157
- batch = Q.shape[:-2]
158
- vec = lambda d: jax.ShapeDtypeStruct(batch + (d,), jnp.float64)
159
- # (x, t, y, z_t, z), then info = [converged, iters].
160
- out_types = [vec(n), vec(p), vec(m), vec(p), vec(p), vec(2)]
161
- call = jax.ffi.ffi_call("elastiqp_solve_warm", out_types, vmap_method="sequential")
162
110
  return call(
163
111
  Q,
164
112
  q,
@@ -173,56 +121,12 @@ def _ffi_solve_warm(
173
121
  eps_abs=np.float64(eps_abs),
174
122
  max_iter=np.int64(max_iter),
175
123
  ruiz=np.int64(bool(ruiz)),
124
+ target_kappa=np.float64(target_kappa),
125
+ warm=np.int64(bool(warm)),
176
126
  method=np.int64(_METHOD_ID[method]),
177
127
  )
178
128
 
179
129
 
180
- @partial(jax.custom_vjp, nondiff_argnums=(10, 11, 12, 13))
181
- def _solve_warm(
182
- Q,
183
- q,
184
- A,
185
- b,
186
- G,
187
- h,
188
- penalty,
189
- x0,
190
- y0,
191
- z0,
192
- eps_abs,
193
- max_iter,
194
- ruiz,
195
- method,
196
- ):
197
- return _ffi_solve_warm(
198
- Q,
199
- q,
200
- A,
201
- b,
202
- G,
203
- h,
204
- penalty,
205
- x0,
206
- y0,
207
- z0,
208
- eps_abs,
209
- max_iter,
210
- ruiz,
211
- method,
212
- )
213
-
214
-
215
- def _solve_warm_fwd(*args):
216
- raise TypeError("elastiqp.jax.solve is not differentiable with warm_start")
217
-
218
-
219
- def _solve_warm_bwd(*args):
220
- raise TypeError("elastiqp.jax.solve is not differentiable with warm_start")
221
-
222
-
223
- _solve_warm.defvjp(_solve_warm_fwd, _solve_warm_bwd)
224
-
225
-
226
130
  def _outer(a, b):
227
131
  return a[..., :, None] * b[..., None, :]
228
132
 
@@ -275,7 +179,9 @@ def _kkt_bwd(res, ct):
275
179
  return Qb, qb, Ab, bb, Gb, hb, penalty_b
276
180
 
277
181
 
278
- @partial(jax.custom_vjp, nondiff_argnums=(7, 8, 9, 10, 11))
182
+ # target_kappa is only used by _solve_fwd; custom_vjp needs the primal and
183
+ # fwd signatures to match. (x0, y0, z0) are ignored unless warm.
184
+ @partial(jax.custom_vjp, nondiff_argnums=(10, 11, 12, 13, 14, 15))
279
185
  def _solve(
280
186
  Q,
281
187
  q,
@@ -284,10 +190,14 @@ def _solve(
284
190
  G,
285
191
  h,
286
192
  penalty,
193
+ x0,
194
+ y0,
195
+ z0,
287
196
  eps_abs,
288
197
  max_iter,
289
198
  ruiz,
290
199
  target_kappa,
200
+ warm,
291
201
  method,
292
202
  ):
293
203
  # Tight solution when not differentiating
@@ -299,10 +209,14 @@ def _solve(
299
209
  G,
300
210
  h,
301
211
  penalty,
212
+ x0,
213
+ y0,
214
+ z0,
302
215
  eps_abs,
303
216
  max_iter,
304
217
  ruiz,
305
218
  0.0,
219
+ warm,
306
220
  method,
307
221
  )
308
222
  return tuple(out[:5]) + (out[10],) # solution + info
@@ -316,10 +230,14 @@ def _solve_fwd(
316
230
  G,
317
231
  h,
318
232
  penalty,
233
+ x0,
234
+ y0,
235
+ z0,
319
236
  eps_abs,
320
237
  max_iter,
321
238
  ruiz,
322
239
  target_kappa,
240
+ warm,
323
241
  method,
324
242
  ):
325
243
  if not target_kappa > 0:
@@ -335,18 +253,27 @@ def _solve_fwd(
335
253
  G,
336
254
  h,
337
255
  penalty,
256
+ x0,
257
+ y0,
258
+ z0,
338
259
  eps_abs,
339
260
  max_iter,
340
261
  ruiz,
341
262
  target_kappa,
263
+ warm,
342
264
  method,
343
265
  )
344
266
  x, t, y, z_t, z, xr, tr, yr, z_t_r, z_r, info = out
345
267
  return (x, t, y, z_t, z, info), (Q, A, G, h, xr, tr, yr, z_t_r, z_r)
346
268
 
347
269
 
348
- def _solve_bwd(eps_abs, max_iter, ruiz, target_kappa, method, res, ct):
349
- return _kkt_bwd(res, ct)
270
+ def _solve_bwd(eps_abs, max_iter, ruiz, target_kappa, warm, method, res, ct):
271
+ # Zero cotangent for the warm start: a converged solve of a QP with a
272
+ # unique solution does not depend on where the solver starts. The
273
+ # (x, y, z) cotangents share the (x0, y0, z0) shapes.
274
+ xb, _, yb, _, zb, _ = ct
275
+ zeros = (jnp.zeros_like(xb), jnp.zeros_like(yb), jnp.zeros_like(zb))
276
+ return _kkt_bwd(res, ct) + zeros
350
277
 
351
278
 
352
279
  _solve.defvjp(_solve_fwd, _solve_bwd)
@@ -396,7 +323,9 @@ def solve(
396
323
  derivatives. Defaults to 1e-3.
397
324
  warm_start (tuple | Result, optional): Explicit warm start, either a
398
325
  previous Result or a (x, y, z) tuple, (das and pdal only).
399
- Defaults to None.
326
+ Differentiable: the gradients are the cold solve's, and the warm
327
+ start itself gets a zero gradient (exact when the QP has a unique
328
+ solution and the solve converges). Defaults to None.
400
329
 
401
330
  Returns:
402
331
  Result: Solution to the elastic QP
@@ -447,32 +376,11 @@ def solve(
447
376
  f"{q.shape}, {b.shape}, {h.shape}; got "
448
377
  f"{x0.shape}, {y0.shape}, {z0.shape}"
449
378
  )
450
- x, t, y, z_t, z, info = _solve_warm(
451
- Q,
452
- q,
453
- A,
454
- b,
455
- G,
456
- h,
457
- penalty,
458
- x0,
459
- y0,
460
- z0,
461
- float(eps_abs),
462
- int(max_iter),
463
- bool(ruiz),
464
- method,
465
- )
466
- return Result(
467
- x=x,
468
- t=t,
469
- y=y,
470
- z_t=z_t,
471
- z=z,
472
- converged=info[..., 0].astype(jnp.int32),
473
- iters=info[..., 1].astype(jnp.int32),
474
- )
475
-
379
+ else:
380
+ # Ignored by the solver (warm=False)
381
+ x0 = jnp.zeros(q.shape, jnp.float64)
382
+ y0 = jnp.zeros(b.shape, jnp.float64)
383
+ z0 = jnp.zeros(h.shape, jnp.float64)
476
384
  x, t, y, z_t, z, info = _solve(
477
385
  Q,
478
386
  q,
@@ -481,10 +389,14 @@ def solve(
481
389
  G,
482
390
  h,
483
391
  penalty,
392
+ x0,
393
+ y0,
394
+ z0,
484
395
  float(eps_abs),
485
396
  int(max_iter),
486
397
  bool(ruiz),
487
398
  float(target_kappa),
399
+ warm_start is not None,
488
400
  method,
489
401
  )
490
402
  # Recall: info = [converged, iters, relax_converged]. The converged
@@ -120,8 +120,8 @@ constexpr Option kOptions[] = {
120
120
  ELASTIQP_D(sing_tol, 0.0,
121
121
  "Working-set LDL' pivot below this = dependent row."),
122
122
  ELASTIQP_D(zero_tol, 0.0,
123
- "Min/max Cholesky pivot ratio of Q_s accepted before adding "
124
- "a prox shift."),
123
+ "Min/max eigenvalue ratio of Q_s (from its Cholesky pivots, "
124
+ "see cond_iters) accepted before adding a prox shift."),
125
125
  ELASTIQP_D(eps_prox, 0.0,
126
126
  "Initial prox shift (x max|diag Q_s|) when Q_s is not PD; 0 "
127
127
  "disables."),
@@ -159,6 +159,10 @@ constexpr Option kOptions[] = {
159
159
  ELASTIQP_I(relax_factor_retries,
160
160
  "relax(): x100 regularization retries on a failed "
161
161
  "factorization."),
162
+ ELASTIQP_I(cond_iters,
163
+ "Inverse-iteration steps confirming the smallest eigenvalue "
164
+ "of Q_s against zero_tol when the pivots look small; 0: "
165
+ "pivots only."),
162
166
  };
163
167
 
164
168
  #undef ELASTIQP_D
@@ -286,7 +286,7 @@ static void options(void) {
286
286
  int i, j, type = -1, ok = 1, unique = 1, n_opt = elastiqp_num_options();
287
287
 
288
288
  elastiqp_setup(&work, &qp, NULL);
289
- check("options: 22 listed", n_opt == 22);
289
+ check("options: 23 listed", n_opt == 23);
290
290
  for (i = 0; i < n_opt; ++i) {
291
291
  ok &= elastiqp_option_info(i, &name, &type, &def, &lo, &hi, &desc) == 0;
292
292
  ok &= name != NULL && desc != NULL && desc[0] != '\0';
@@ -215,6 +215,7 @@ void SettingsTable() {
215
215
  FIELD(refactor_tol);
216
216
  FIELD(relax_reg);
217
217
  FIELD(relax_factor_retries);
218
+ FIELD(cond_iters);
218
219
  #undef FIELD
219
220
  elastiqp_free(work);
220
221
  Check("option table covers every settings field",
@@ -14,9 +14,10 @@
14
14
  // 3. Robustness: a badly scaled parametrization with and without Ruiz;
15
15
  // dependent / inconsistent equalities (certificate before iterating);
16
16
  // every inequality row duplicated (dependent working sets, the cycle
17
- // guard's regime, refactors()); the drift-gated Ruiz refresh through a
18
- // row exploding past the limit_scaling clamp and a row decaying to the
19
- // noise floor.
17
+ // guard's regime, refactors()); a near-singular Q the Cholesky pivots
18
+ // hide (the eigenvalue check, and the prox fallback without it); the
19
+ // drift-gated Ruiz refresh through a row exploding past the
20
+ // limit_scaling clamp and a row decaying to the noise floor.
20
21
  // 4. Factorization reuse: constant Q, a fifth of G changing per tick,
21
22
  // same answers as a full refactorization and only those rows re-solved.
22
23
  // 5. The pinch/release ticks of test_gap_creep: saturated -> active ->
@@ -294,6 +295,58 @@ void ProxLoop(std::mt19937& rng) {
294
295
  "status");
295
296
  }
296
297
 
298
+ // Near-singular Q whose null direction is small on the trailing coordinates
299
+ // (as in diff IK, J'J + 1e-12 I with 6 task rows on 7 DOF): the Cholesky
300
+ // pivots overstate its smallest eigenvalue, so only the inverse-iteration
301
+ // check turns the prox shift on up front. Without it (cond_iters = 0), the
302
+ // inner loop can break down and must recover through the prox fallback.
303
+ void NearSingularQ(std::mt19937& rng) {
304
+ const int n = 7, p = 14, cases = 20;
305
+ int solved = 0, prox = 0, escalated = 0;
306
+ int solved_off = 0, escalated_off = 0;
307
+ double worst = 0.0;
308
+ for (int k = 0; k < cases; ++k) {
309
+ VectorXd v = problem_gen::Randn(rng, n, 1);
310
+ v.tail(2) *= 0.05;
311
+ v.normalize();
312
+ const MatrixXd J = problem_gen::Randn(rng, n - 1, n) *
313
+ (MatrixXd::Identity(n, n) - v * v.transpose());
314
+ const MatrixXd Q = J.transpose() * J + 1e-12 * MatrixXd::Identity(n, n);
315
+ const VectorXd q = -J.transpose() * problem_gen::Randn(rng, n - 1, 1);
316
+ MatrixXd G(p, n);
317
+ G << MatrixXd::Identity(n, n), -MatrixXd::Identity(n, n);
318
+ const VectorXd h = VectorXd::Constant(p, 0.5) +
319
+ 0.1 * problem_gen::Randn(rng, p, 1);
320
+ const VectorXd w = VectorXd::Constant(p, 1e3);
321
+ das::Solver s;
322
+ s.setup(Q, q, G, h, w);
323
+ const elastiqp::Solution sol = s.solve();
324
+ solved += sol.converged;
325
+ prox += s.proximal();
326
+ escalated += s.prox_escalated();
327
+ QPData qp;
328
+ qp.Q = Q;
329
+ qp.q = q;
330
+ qp.A = MatrixXd(0, n);
331
+ qp.b = VectorXd(0);
332
+ qp.G = G;
333
+ qp.h = h;
334
+ worst = std::max(worst, Kkt(qp, w, sol));
335
+ das::Solver off;
336
+ off.settings.cond_iters = 0;
337
+ off.setup(Q, q, G, h, w);
338
+ solved_off += off.solve().converged;
339
+ escalated_off += off.prox_escalated();
340
+ }
341
+ Check("near-singular Q: solved, prox on up front",
342
+ solved == cases && prox == cases && escalated == 0,
343
+ static_cast<double>(solved), "solved");
344
+ Check("near-singular Q: elastic KKT residual", worst <= 1e-6, worst, "kkt");
345
+ Check("near-singular Q, cond_iters = 0: solved via the prox fallback",
346
+ solved_off == cases && escalated_off > 0,
347
+ static_cast<double>(escalated_off), "escalated");
348
+ }
349
+
297
350
  // Drift-gated Ruiz refresh through the update path: a row exploding past
298
351
  // the 1e4 clamp in one tick (rescaled() fires, drift back to O(1)) and a
299
352
  // row decaying to the noise floor (left alone, solved as if absent).
@@ -470,6 +523,7 @@ void RandomSuite() {
470
523
  EqualityCases(rng);
471
524
  DegenerateCase(rng);
472
525
  ProxLoop(rng);
526
+ NearSingularQ(rng);
473
527
  RuizRefresh(rng);
474
528
  ReuseBench();
475
529
  }
@@ -14,6 +14,29 @@ import numpy as np
14
14
  RESULTS = []
15
15
 
16
16
 
17
+ def ffi_solve_cold(
18
+ Q, q, A, b, G, h, penalty, eps_abs, max_iter, ruiz, target_kappa, method
19
+ ):
20
+ # Raw FFI call without a warm start (x0, y0, z0 are ignored placeholders)
21
+ zeros = (jnp.zeros_like(q), jnp.zeros_like(b), jnp.zeros_like(h))
22
+ return elastiqp.jax._ffi_solve(
23
+ Q,
24
+ q,
25
+ A,
26
+ b,
27
+ G,
28
+ h,
29
+ penalty,
30
+ *zeros,
31
+ eps_abs,
32
+ max_iter,
33
+ ruiz,
34
+ target_kappa,
35
+ False,
36
+ method,
37
+ )
38
+
39
+
17
40
  def check(name, ok, detail=""):
18
41
  RESULTS.append(ok)
19
42
  print(f" {name:<52s} {detail:<24s} {'OK' if ok else 'FAIL'}")
@@ -198,7 +221,7 @@ def main():
198
221
  return w_loss @ s.x + w_t @ s.t
199
222
 
200
223
  def loss_relaxed(Q_, q_, A_, b_, G_, h_, penalty_, kap=kappa, method="das"):
201
- out = elastiqp.jax._ffi_solve(
224
+ out = ffi_solve_cold(
202
225
  Q_,
203
226
  q_,
204
227
  A_,
@@ -317,7 +340,7 @@ def main():
317
340
  # relaxation would otherwise mean silently wrong gradients). s_t = t and
318
341
  # s_in = h + t - Gx are the slacks of t >= 0 and Gx - t <= h; s_in is
319
342
  # reconstructed from x, so it carries the O(tol) primal residual.
320
- out = elastiqp.jax._ffi_solve(*args, 1e-11, 300, False, kappa, "das")
343
+ out = ffi_solve_cold(*args, 1e-11, 300, False, kappa, "das")
321
344
  xr, tr, z_t_r, z_r, info = out[5], out[6], out[8], out[9], out[10]
322
345
  s_in_r = args[5] + tr - args[4] @ xr
323
346
  comp = max(
@@ -334,7 +357,7 @@ def main():
334
357
  # silently: an absurd kappa converges the solve (info[0]) but stalls the
335
358
  # relaxation (info[2]), and solve() folds that into converged on the
336
359
  # differentiated path.
337
- out = elastiqp.jax._ffi_solve(*args, 1e-11, 300, False, 1e8, "das")
360
+ out = ffi_solve_cold(*args, 1e-11, 300, False, 1e8, "das")
338
361
  tight_ok, relax_bad = float(out[10][0]) == 1.0, float(out[10][2]) == 0.0
339
362
 
340
363
  def conv_smooth(q_):
@@ -468,9 +491,7 @@ def main():
468
491
  try:
469
492
  jax.grad(
470
493
  lambda q_: jnp.sum(
471
- elastiqp.jax.solve(
472
- Qx, q_, Gx, hx, 10.0, A=Ax, b=bx, target_kappa=0.0
473
- ).x
494
+ elastiqp.jax.solve(Qx, q_, Gx, hx, 10.0, A=Ax, b=bx, target_kappa=0.0).x
474
495
  )
475
496
  )(qx)
476
497
  msg = None
@@ -517,7 +538,7 @@ def main():
517
538
  for m in elastiqp.jax.METHODS
518
539
  }
519
540
  relaxed = {
520
- m: elastiqp.jax._ffi_solve(*args, 1e-11, 300, False, kappa, m)
541
+ m: ffi_solve_cold(*args, 1e-11, 300, False, kappa, m)
521
542
  for m in elastiqp.jax.METHODS
522
543
  }
523
544
  for m in ("pdal", "ipm"):
@@ -566,9 +587,9 @@ def main():
566
587
  int(a_sol.converged) == 1 and da < FFI_NB_TOL,
567
588
  f"|dx|={da:.1e}",
568
589
  )
569
- a_jit = jax.jit(
570
- lambda q_: elastiqp.jax.solve(Qx, q_, Gx, hx, 10.0, A=Ax, b=bx).x
571
- )(qx)
590
+ a_jit = jax.jit(lambda q_: elastiqp.jax.solve(Qx, q_, Gx, hx, 10.0, A=Ax, b=bx).x)(
591
+ qx
592
+ )
572
593
  check("das under jit", float(jnp.abs(a_jit - a_sol.x).max()) == 0.0, "")
573
594
  # Hard rows (penalty = inf, das only): finite gradients, none w.r.t. the
574
595
  # hard row's penalty, close to a stiff elastic penalty in its place.
@@ -671,8 +692,15 @@ def main():
671
692
  check("warm sequence has working-set changes", ws_changed > 0, f"{ws_changed}")
672
693
  try:
673
694
  elastiqp.jax.solve(
674
- Qw, drift[0][0], Gw, drift[0][1], 10.0, A=Aw, b=drift[0][2],
675
- method="ipm", warm_start=prev,
695
+ Qw,
696
+ drift[0][0],
697
+ Gw,
698
+ drift[0][1],
699
+ 10.0,
700
+ A=Aw,
701
+ b=drift[0][2],
702
+ method="ipm",
703
+ warm_start=prev,
676
704
  )
677
705
  ipm_raised = False
678
706
  except ValueError:
@@ -755,22 +783,57 @@ def main():
755
783
  f"|dx|={dxv:.1e} max iters={int(v_warm.iters.max())}",
756
784
  )
757
785
 
758
- try:
759
- jax.grad(
760
- lambda q_: jnp.sum(
761
- elastiqp.jax.solve(
762
- Qw, q_, Gw, h_k, 10.0, A=Aw, b=b_k, warm_start=p_sol
763
- ).x
786
+ # A warm-started solve has the cold solve's gradients (the solution does
787
+ # not depend on the start), and a zero gradient for the warm start itself
788
+ def loss_w(Q_, q_, G_, h_, pen_, warm, method):
789
+ s = elastiqp.jax.solve(
790
+ Q_, q_, G_, h_, pen_, A=Aw, b=b_k, method=method, warm_start=warm
791
+ )
792
+ return jnp.sum(s.x**2) + jnp.sum(jnp.sin(s.x))
793
+
794
+ pen_k = jnp.full(h_k.shape, 10.0)
795
+ warm = (p_sol.x, p_sol.y, p_sol.z)
796
+ for method in ("das", "pdal"):
797
+ g_cold = jax.grad(loss_w, argnums=(0, 1, 2, 3, 4))(
798
+ Qw, q_k, Gw, h_k, pen_k, None, method
799
+ )
800
+ *g_warm, g_start = jax.grad(loss_w, argnums=(0, 1, 2, 3, 4, 5))(
801
+ Qw, q_k, Gw, h_k, pen_k, warm, method
802
+ )
803
+ dg = max(
804
+ float(jnp.abs(a - b).max()) for a, b in zip(g_cold, g_warm, strict=True)
805
+ )
806
+ g0 = max(float(jnp.abs(g).max()) for g in g_start)
807
+ check(
808
+ f"grad with warm_start ({method}) == cold",
809
+ dg < 1e-8 and g0 == 0.0,
810
+ f"|dgrad|={dg:.1e} |grad warm start|={g0:.1e}",
811
+ )
812
+
813
+ # Warm-started lax.scan works under jit + grad and matches the cold scan.
814
+ # The carry's cotangent is zero by construction, so this is a smoke test
815
+ # of the warm path, not an independent check of the gradient through it
816
+ def scan_loss(q_, warm_carry):
817
+ def body(carry, h_):
818
+ s = elastiqp.jax.solve(
819
+ Qw,
820
+ q_,
821
+ Gw,
822
+ h_,
823
+ 10.0,
824
+ A=Aw,
825
+ b=b_k,
826
+ warm_start=carry if warm_carry else None,
764
827
  )
765
- )(q_k)
766
- msg = None
767
- except TypeError as e:
768
- msg = str(e)
769
- check(
770
- "grad with warm_start raises",
771
- msg is not None and "warm_start" in msg and "not differentiable" in msg,
772
- "",
773
- )
828
+ return (s.x, s.y, s.z), jnp.sum(s.x**2)
829
+
830
+ _, losses = jax.lax.scan(body, warm, hs)
831
+ return jnp.sum(losses)
832
+
833
+ g_cold = jax.jit(jax.grad(scan_loss), static_argnums=1)(q_k, False)
834
+ g_warm = jax.jit(jax.grad(scan_loss), static_argnums=1)(q_k, True)
835
+ dgs = float(jnp.abs(g_cold - g_warm).max())
836
+ check("grad through a warm-started lax.scan", dgs < 1e-8, f"|dgrad|={dgs:.1e}")
774
837
  try:
775
838
  elastiqp.jax.solve(
776
839
  Qw,
@@ -825,8 +888,15 @@ def main():
825
888
 
826
889
  def x0_of_h0(h0, method):
827
890
  return elastiqp.jax.solve(
828
- Qb, qb, Gb, hb.at[0].set(h0), 10.0, A=Ab, b=bb,
829
- method=method, target_kappa=1e-9,
891
+ Qb,
892
+ qb,
893
+ Gb,
894
+ hb.at[0].set(h0),
895
+ 10.0,
896
+ A=Ab,
897
+ b=bb,
898
+ method=method,
899
+ target_kappa=1e-9,
830
900
  ).x[0]
831
901
 
832
902
  for method in ["das", "pdal", "ipm"]:
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes