Skip to content

Sampling book update and clean up#80

Merged
junpenglao merged 4 commits intomainfrom
phase1-jax-api-fixes
Mar 26, 2026
Merged

Sampling book update and clean up#80
junpenglao merged 4 commits intomainfrom
phase1-jax-api-fixes

Conversation

@junpenglao
Copy link
Copy Markdown
Member

Co-Authored-By: Claude Sonnet 4.6 noreply@anthropic.com

junpenglao and others added 4 commits March 26, 2026 09:52
Replaces jax.tree_util.tree_map/flatten and jax.tree_map with the
idiomatic jax.tree.map/flatten introduced in JAX 0.9.x. Also removes
stale reference to jaxopt in pathfinder.md comment (replaced by
optax.lbfgs in BlackJAX #817).

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
blackjax.pathfinder(logdensity_fn) now returns a VIAlgorithm whose
init/sample interface replaces the old approximate/sample interface.
Update code and prose accordingly.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
sparse_regression.md:
- Replace aesara/aeppl model definition with pure JAX log-density.
  aesara uses np.obj2sctype removed in NumPy 2.0 and the ecosystem
  is no longer maintained.
- Fix deprecated pandas delim_whitespace=True -> sep=r"\s+"

probabilistic_ode_solver_parameter_estimation.md:
- Add jax.tree_map compatibility shim; probdiffeq 0.1.x uses the
  top-level jax.tree_map shorthand removed in JAX 0.9.0.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
The 0.1.x API (solution_routines, solvers, implementations.recipes,
strategies.filters) was completely replaced in 0.2+. New API uses
ivpsolve / ivpsolvers / stats / taylor modules.

Key changes:
- solver built via prior_wiener_integrated + correction_ts0 + strategy_smoother
- solve_fixed_grid / solve_adaptive_save_at replace old solution_routines
- stats.log_marginal_likelihood_terminal_values replaces
  condition_on_qoi_observation + marginals.logpdf
- ssm built once outside JIT; init_cond rebuilt per theta
- plot_solution updated for sol.u list-of-coefficients layout
- vmap solution_samples iterated by index (batched IVPSolution)
- Drop probdiffeq<0.2.0 pin in requirements.txt

Figure outputs verified within noise of MCMC stochasticity.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
@junpenglao junpenglao merged commit a771b1d into main Mar 26, 2026
1 check passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant