blackjax-devs / blackjax-devs/tuningfork
bug: multi-chain laplace warmup has_aux TypeError (vmap+grad over LaplaceMarginal)
- Dominant language
- Python
- Stars
- 3
- Forks
- 0
- Avg merge
- 7h 24m
- Merged PRs (30d)
- 4
Description
## Summary
When running `emit_low_recipe_for_cell` (or any path that vmaps the laplace warmup across chains with `num_chains > 1`), the warmup fails immediately with a JAX `has_aux` TypeError:
```
expected two-element tuple from function with has_aux=True,
but got LinearizeTracer / GradTracer
```
## Root cause
`blackjax.window_adaptation` with a laplace inner kernel calls `jax.value_and_grad(laplace, has_aux=True)(phi)` internally. The `LaplaceMarginal.__call__` returns `(lp, theta_star)` — a `(float, pytree)` pair satisfying `has_aux=True`. However, when the warmup is **vmapped across chains** (`jax.vmap(warmup_run)(...)`), the traced function gets wrapped in a `BatchTracer` and the `(lp, theta_star)` output gets linearized, producing a `LinearizeTracer` instead of the expected tuple at JAX's internal `has_aux` demultiplexing point.
## Workaround
Use `warmup_num_chains=[1, 1]` on the recipe (run warmup on a SINGLE chain, then broadcast to `num_chains` sampling chains). This avoids the vmap of the laplace warmup entirely. Already stamped on the gp_regression HIGH recipe (`warmup_num_chains=[1,1]` per PR #72 findings).
## Reproduction
1. Call `emit_low_recipe_for_cell("gp_regression", "window_adaptation_dense_imm", "laplace_mhmc", num_chains=4, ...)` without the `warmup_num_chains=[1,1]` recipe stamp.
2. Observe the `has_aux` TypeError in the first warmup step.
## Notes
- The bug is in the blackjax `window_adaptation` + laplace-kernel combination — `window_adaptation` internally uses `jax.value_and_grad(..., has_aux=True)` and this doesn't compose cleanly with `vmap` for the LaplaceMarginal output structure.
- The `run_recipe_to_idata` path with `warmup_num_chains=[1,1]` is the **only** known-working multi-chain recert path for gp_regression × laplace_mhmc.
- This likely affects all laplace-family recipes for models with D > a few tens (where per-step cost makes multi-chain warmup slow anyway).
## Relevant PRs
- PR #72: discovered the original vmap warmup deadlock (different symptom, same root) — fixed by `warmup_num_chains=[1,1]`
- PR #81: re-cert attempt that triggered this error path; sidestepped via single-chain warmup
## Priority
P2 — sidestepped in all current recipes. Worth fixing upstream in blackjax `window_adaptation` to support `has_aux=True` laplace logdensity functions under vmap.
Contributor guide
Assessment
This issue has not been assessed yet.