[pratik-ingle]

Problem

1. linalg.solve is annotated A: Array | ProcessedDriftMatrix, but passing a ProcessedDriftMatrix fails, because the OU mean is computed with jnp.linalg.solve(A, b):

   A = jnp.array([[3, 2], [2, 4.0]])
   thermox.linalg.solve(thermox.utils.preprocess_drift_matrix(A), jnp.array([1, 2.0]))
   # TypeError: jnp.linalg.solve requires ndarray or scalar arguments,
   # got <class 'thermox.utils.ProcessedDriftMatrix'> at position 0.

   Passing a preprocessed matrix matters when many right-hand sides share one A (e.g. a multi-column Kalman gain), since preprocessing is the O(d³) step.

2. The process always starts at zero, so for small num_samples the time average is biased toward zero. On the matrix from test_linear_system (dt=1, averaged over 500 keys), N=10 gives a 9.9% relative bias; starting at the solution gives 0.5%, which is the sampling noise. The bias grows with conditioning (random SPD with κ=10: 50% at N=10, 7.8% at N=100). burnin also reduces it but discards samples; when a good guess exists (a nearby system's solution, or the previous iterate in an iterative-refinement loop), starting there removes it at no cost.

Changes

• If A is a ProcessedDriftMatrix, the mean is computed from its stored eigendecomposition (O(d²), no extra factorization). Array inputs behave exactly as before.
• New keyword argument x0 (default None, meaning zeros as before), added after associative_scan so positional calls are unaffected.
• Docstring updated; two tests added to tests/test_linalg.py.

Context

Found while benchmarking thermox on robotics workloads: https://github.com/pratik-ingle/thermox-robot-bench

[Copilot]

🟢 Approval recommended

The changes preserve existing behavior, match the sampler contract, and have no identified blocking issues.

No open issues 🎉

What changed in this PR

Updates linalg.solve to support preprocessed drift matrices and configurable initial states while preserving existing defaults.

Changes:

• Reuses stored eigendecompositions for preprocessed inputs.
• Adds optional x0 without shifting existing positional arguments.
• Documents both features and adds regression tests.

File	Description
thermox/​linalg.py	Supports preprocessed matrices and optional initial states.
tests/​test_linalg.py	Tests preprocessed-input equivalence and initial-state bias reduction.

_{Review effort: Balanced}

[SamDuffield]

Nice catch! Just for your reference, the linalg suite is more for fun and the core thermox is the rest of the package which is more flexible