bug: ConditionedINR + Bayesian conditioners broken — builder signature mismatch, and generated siblings need unique scopes
- Dominant language
- Python
- Stars
- 1
- Forks
- 0
- Avg merge
- 19h 2m
- Merged PRs (30d)
- 16
Description
## Problem
The documented Bayesian FiLM workflow — `ConditionedINR.init(inner, conditioner_cls=BayesianAffineModulation, ...)` as shown in `docs/notebooks/conditioning.ipynb` (cell 23) — is broken by **two stacked defects**:
**1. Pre-existing (broke with the geonnax 0.0.5 migration, independent of #187):** geonnax's `_build_conditioner` passes `key=` to `conditioner_cls.init`, but `BayesianAffineModulation.init` (current signature: `(num_features, cond_dim, *, gamma_activation, prior_std, pyrox_name)`) does not accept `key`. The exact tutorial call raises `TypeError: BayesianAffineModulation.init() got an unexpected keyword argument 'key'` **on main today**. The notebook's stored outputs predate the migration.
**2. Introduced by #187 (latent until defect 1 is fixed):** `ConditionedINR.init` builds one conditioner per non-readout layer from a single shared `conditioner_kwargs` — multiple *unnamed* instances of one class. The tutorial's stored output shows they used to disambiguate via the old id-scoped fallback (`BayesianAffineModulation_73cf38418cd0.gen_W`, `…8f50.gen_W`, `…a2510.gen_W`). Under the deterministic class-name fallback they would all share `BayesianAffineModulation.*` and the duplicate-site guard (correctly) raises on the second conditioner. Note that passing `pyrox_name` via `conditioner_kwargs` cannot help: the kwargs are shared, so every generated conditioner would get the *same* name.
## Reproduction
```python
import jax.numpy as jnp, jax.random as jr
from numpyro import handlers
from pyrox_nn import SIREN, BayesianAffineModulation, ConditionedINR
inner = SIREN.init(1, 32, 1, depth=4, key=jr.key(0))
wrapped = ConditionedINR.init( # TypeError today (defect 1)
inner, conditioner_cls=BayesianAffineModulation, cond_dim=8, key=jr.key(1)
)
with handlers.seed(rng_seed=0), handlers.trace():
wrapped(jnp.zeros((4, 1)), jnp.zeros((4, 8))) # would hit defect 2 once 1 is fixed
```
## Expected Behavior
The Bayesian FiLM composite constructs and traces with one distinct, **deterministic** site scope per generated conditioner (stable across `jit`/`tree_at`/checkpoint reloads, per the #184/#187 naming contract).
## Environment
- pyrox-nn 0.1.0 (post-#187 branch, but defect 1 reproduces on main @ b0ef508 too)
- geonnax 0.0.5, numpyro 0.21.0, jax 0.10.2, equinox 0.13.8, Python 3.13
## References & Existing Code
- Generated-conditioner loop: `geonnax` `ConditionedINR.init` — `for i, k in enumerate(keys): conditioners.append(_build_conditioner(conditioner_cls, ..., **conditioner_kwargs))`
- Bayesian conditioners: `packages/pyrox-nn/src/pyrox_nn/_conditioning.py` (`BayesianAffineModulation`, `BayesianConcatConditioner`, `BayesianHyperLinear` — all `init` classmethods lack `key` compat)
- Tutorial: `docs/notebooks/conditioning.ipynb` cell 23 (stored outputs show the old id-scoped names)
- Naming contract this must respect: #184 / #187 (`_pyrox_scope_name` class-name fallback; duplicate-site guard)
## Proposed Fix Sketch
1. **Builder compat:** add `key: PRNGKeyArray | None = None` (accepted, unused — Bayesian conditioners sample at trace time) to the Bayesian conditioner `init` classmethods so geonnax's `_build_conditioner` protocol works.
2. **Unique generated names:** when `pyrox_name is None` in these `init` factories, assign a per-class monotonic counter name (e.g. `BayesianAffineModulation_0`, `_1`, …) **stored in the static `pyrox_name` field at construction**. Unlike the removed id fallback, the name is baked into the pytree, so it survives `tree_at`/flatten/checkpoints; it is deterministic for a fixed construction order (same script → same names). Document that cross-process stability requires identical construction order, or an explicit name.
- Alternative considered: a pyrox-side `ConditionedINR` wrapper that names conditioners `{base}.layer{i}` — better names, but reimplements geonnax's builder loop (coupling to its internals).
3. **Tutorial:** re-execute the Bayesian FiLM cells of `conditioning.ipynb` once fixed (stored outputs currently show the defunct id-scoped names).
## Testing
- [ ] `packages/pyrox-nn/tests/nn/test_conditioning.py::test_conditioned_inr_bayesian_film_traces` — the repro above constructs, traces cleanly, and registers `depth-1` distinct `gen_W`/`gen_b` pairs
- [ ] Generated names stable across `jax.tree.flatten`/`unflatten` of the wrapped model
- [ ] Explicit `pyrox_name` in `conditioner_kwargs` raises a clear error or is documented as unsupported (shared kwargs → guaranteed collision)
## Relationships
- Parent (theme epic, if any): #
- Blocked by: #
- Blocks: #
Found via Codex review round 8 on #187 (thread: https://github.com/jejjohnson/pyrox/pull/187#discussion_r3608131784 area); triaged out of #187 because defect 1 predates it and the counter-naming scheme is a naming-contract decision deserving its own review.
Contributor guide
Research direction
Start with ConditionedINR.init and _build_conditioner, then inspect packages/pyrox-nn/src/pyrox_nn/_conditioning.py for the Bayesian conditioner init methods. Run packages/pyrox-nn/tests/nn/test_conditioning.py::test_conditioned_inr_bayesian_film_traces and verify construction, tracing, distinct generated sites, and name stability through tree flattening. Re-execute the Bayesian FiLM cells in docs/notebooks/conditioning.ipynb after the tests pass.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- machine-learning
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Quiet
- Clarity
- Mostly clear
- Newbie friendliness
- 52/100