Of course, after I posted this I found another question addressing kinda the same issue. Translating the solution given there to the even simpler setting I have here gives
def model():
# draw a decision
A = numpyro.sample(
"A",
dist.Bernoulli(0.5),
infer={"enumerate": "parallel"}
)
print(f"A.shape={A.shape}")
# evaluate likelihood given the selected params
with numpyro.plate("data_plate", N) as i:
# choose (using Vindex for correctness) the corresponding set of params
chosen_p_B = Vindex(all_p_B)[A, i]
numpyro.deterministic('chosen_p_B', chosen_p_B)
print(f"chosen_p_B.shape={chosen_p_B.shape}")
B = numpyro.sample(
"B",
dist.Bernoulli(chosen_p_B),
obs = jnp.ones(N, dtype=A.dtype)
)
Prior simulation works
A.shape=()
i.shape=(5,)
chosen_p_B.shape=(5,)
>>> for k,v in exec_trace.items():
... if v['type'] != 'plate':
... print(f"{k}: {v['value']}")
A: 0
chosen_p_B: [0.1 0.2 0.3 0.4 0.5]
B: [1 1 1 1 1]
as well as initialization
>>> (
... params_info,
... potential_fn_gen,
... postprocess_fn,
... model_trace,
... ) = numpyro.infer.util.initialize_model(
... init_key,
... model,
... dynamic_args=True,
... init_strategy=numpyro.infer.init_to_sample, # initialize from the prior
... )
A.shape=()
i.shape=(5,)
chosen_p_B.shape=(5,)
A.shape=()
i.shape=(5,)
chosen_p_B.shape=(5,)
A.shape=(2, 1)
i.shape=(5,)
chosen_p_B.shape=(2, 5)
I guess I don’t really understand what’s going, but the practical effect is clear: the singleton dimension in chosen_p_B is not created.
Quick check that the implementation is correct
# P(obs) = sum_k p_k P(y|k) = p sum_k P(y|k)
# => -logP(obs) = -log(p) - logsumexp(logP(y|k))
init_params = params_info[0]
potential_fn = potential_fn_gen()
conditional_logliks = jnp.log(all_p_B).sum(axis=1)
logp = jnp.log(0.5)
assert jnp.allclose(
potential_fn(init_params),
-logp-jnp.logaddexp(conditional_logliks[0], conditional_logliks[1])
)