blackjax-devs / blackjax-devs/tuningfork

bug: multi-chain laplace warmup has_aux TypeError (vmap+grad over LaplaceMarginal)

Open
#82 1 comment 0 reactions 0 assignees View on GitHub
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

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.