chore(engine): remove the dead, broken JAX MC backend#99
Conversation
The engine's JAX Monte-Carlo backend was dead and doubly broken: - solve_mc_jax built a batched JaxParams with only 18 of the dataclass's 20 required fields (missing node_fu_correction_applicable, drug_fu_correction_liver) → TypeError on construction. - jax / diffrax are not project dependencies (absent from pyproject and requirements-lock), so importing solver_jax raised ImportError first. - No caller ever selected backend="jax" (production defaults to scipy); every JAX test skipped via importorskip since jax isn't installed. It was labelled "UDE roadmap Phase 0", but UDE is an enumerated dead-end, so the scaffold has no live consumer. Removed the whole engine JAX subsystem: - delete engine/solver_jax.py, params_jax.py, rhs_jax.py - drop the backend=="jax" branch in uncertainty.py (scipy + surrogate remain) - delete test_solver_jax.py, test_jax_scipy_parity.py, test_rhs_jax_unsupported_flux.py - remove the two skipped JAX-vs-SciPy parity tests in test_active_transport*.py The surrogate backend (a separate subsystem) is untouched. SciPy is and remains the only production MC backend. tests/unit/test_active_transport*.py + test_uncertainty.py pass (15).
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 33c59ed02d
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| @@ -1,363 +0,0 @@ | |||
| """JAX/Diffrax-based ODE solver — differentiable alternative to SciPy LSODA. | |||
There was a problem hiding this comment.
In this revision, running the tracked CLI PYTHONPATH=src python scripts/probe_oatp_phase2a_per_drug.py pravastatin 0.1 fails immediately because scripts/probe_oatp_phase2a_per_drug.py:26 still imports sisyphus.engine.solver_jax. Since this script is the per-drug Phase 2A stiff-ODE diagnostic, deleting the module without switching the probe to SciPy or removing the script leaves a documented tool broken.
Useful? React with 👍 / 👎.
| (V2-compatible). | ||
| backend: Solver backend to use (``"scipy"``, ``"jax"``, or | ||
| ``"surrogate"``). | ||
| backend: Solver backend to use (``"scipy"`` or ``"surrogate"``). |
There was a problem hiding this comment.
When an existing caller or experiment config still passes backend="jax", this change no longer takes the JAX path or raises; after the surrogate branch it falls through into the SciPy loop. That silently changes the requested backend and can mislabel performance/results comparisons, so unsupported backend values should be rejected explicitly before the SciPy default path.
Useful? React with 👍 / 👎.
1 participant
The engine's JAX Monte-Carlo backend was dead and doubly broken — verified before removal:
solve_mc_jaxbuilt a batchedJaxParamswith 18 of the 20 required fields (missingnode_fu_correction_applicable,drug_fu_correction_liver) →TypeErroron construction.jax/diffraxare not project dependencies (absent frompyproject.toml+requirements-lock.txt), so importingsolver_jaxraisedImportErrorfirst.backend="jax"(production defaults toscipy); every JAX test skipped viaimportorskipsince jax isn't installed.Labelled "UDE roadmap Phase 0", but UDE is an enumerated dead-end, so the scaffold has no live consumer.
Removed (−1539 lines)
engine/solver_jax.py,params_jax.py,rhs_jax.pybackend=="jax"branch inuncertainty.py(scipy+surrogatebackends remain; docstrings updated)test_solver_jax.py,test_jax_scipy_parity.py,test_rhs_jax_unsupported_flux.pytest_active_transport*.pyThe surrogate backend (a separate subsystem, also imports jax) is intentionally untouched — out of scope. SciPy is and remains the only production MC backend.
Verification
test_active_transport.py+test_active_transport_direction.py+test_uncertainty.py→ 15 passed; ruff clean; no residual references to the deleted modules.