Enumeration and conditionally independent observations

numpyro
1

Hi – I’m working on a complex model that requires enumeration combined with a likelihood defined by conditionally independent outcomes, given some parameters.

I’ve been able to distill the issue by taking the Toy Mixture Model example and adapting it to my situation as best I can

import jax
from jax import numpy as jnp

import numpyro
from numpyro import distributions as dist
from numpyro.ops.indexing import Vindex

Start by defining some fixed parameters

N=5
all_p_B = jnp.stack((
    jnp.arange(1, N+1) / (2*N),
    jnp.arange(N + 1,2*N + 1)/(2*N)
))

>>> print(all_p_B)
[[0.1        0.2        0.3        0.4        0.5       ]
 [0.6        0.7        0.8        0.90000004 1.        ]]

The models is then

def model():
    # draw a decision
    A = numpyro.sample(
        "A",
        dist.Bernoulli(0.5),
        infer={"enumerate": "parallel"}
    )
    print(f"A.shape={A.shape}")

    # choose (using Vindex for correctness) the corresponding set of params
    chosen_p_B = Vindex(all_p_B)[A]
    numpyro.deterministic('chosen_p_B', chosen_p_B)
    print(f"chosen_p_B.shape={chosen_p_B.shape}")

    # evaluate likelihood given the selected params for each observation
    with numpyro.plate("data_plate", N):
        B = numpyro.sample(
            "B",
            dist.Bernoulli(chosen_p_B),
            obs = jnp.ones(N, dtype=A.dtype)
        )

Now, sampling from the prior works

rng_key = jax.random.key(5)
rng_key, seed_key = jax.random.split(rng_key)
seeded_model = numpyro.handlers.seed(model, seed_key)
exec_trace = numpyro.handlers.trace(seeded_model).get_trace()

A.shape=()
chosen_p_B.shape=(5,)

>>> for k,v in exec_trace.items():
...     if v['type'] != 'plate':
...         print(f"{k}: {v['value']}")
... 
A: 1
chosen_p_B: [0.6        0.7        0.8        0.90000004 1.        ]
B: [1 1 1 1 1]

But initialization fails

rng_key, init_key = jax.random.split(rng_key)
(
    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=()
chosen_p_B.shape=(5,)
A.shape=()
chosen_p_B.shape=(5,)
A.shape=(2, 1)
chosen_p_B.shape=(2, 1, 5)
Traceback (most recent call last):
  File "<stdin>", line 6, in <module>
  File "[...]/.venv/lib/python3.12/site-packages/numpyro/infer/util.py", line 751, in initialize_model
    (init_params, pe, grad), is_valid = find_valid_initial_params(
                                        ^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "[...]/.venv/lib/python3.12/site-packages/numpyro/infer/util.py", line 473, in find_valid_initial_params
    (init_params, pe, z_grad), is_valid = _find_valid_params(
                                          ^^^^^^^^^^^^^^^^^^^
  File "[...]/.venv/lib/python3.12/site-packages/numpyro/infer/util.py", line 459, in _find_valid_params
    _, _, (init_params, pe, z_grad), is_valid = init_state = body_fn(init_state)
                                                             ^^^^^^^^^^^^^^^^^^^
  File "[...]/.venv/lib/python3.12/site-packages/numpyro/infer/util.py", line 443, in body_fn
    pe, z_grad = value_and_grad(potential_fn)(params)
                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "[...]/.venv/lib/python3.12/site-packages/numpyro/infer/util.py", line 325, in potential_energy
    log_joint, model_trace = log_density_(
                             ^^^^^^^^^^^^^
  File "[...]/.venv/lib/python3.12/site-packages/numpyro/contrib/funsor/infer_util.py", line 329, in log_density
    result, model_trace, _ = _enum_log_density(
                             ^^^^^^^^^^^^^^^^^^
  File "[...]/.venv/lib/python3.12/site-packages/numpyro/contrib/funsor/infer_util.py", line 298, in _enum_log_density
    raise ValueError(
ValueError: Expected the joint log density is a scalar, but got (2,). There seems to be something wrong at the following sites: {'_pyro_dim_3'}.

The error message mentions this _pyro_dim_3 site that seems to be created by the enumeration procedure, so it doesn’t give me much information on how to proceed.

I’m sure there is something really obvious that I’m missing but I just can’t figure it out. Any pointers would be really appreciated.

2

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])
)