jejjohnson / jejjohnson/pyrox

bug: ConditionedINR + Bayesian conditioners broken — builder signature mismatch, and generated siblings need unique scopes

Open
#188 0 comments 0 reactions 0 assignees View on GitHub
area:code layer:1-components priority:p2 type:bug
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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.