Parallel enumeration: wrong posterior weights (or _pyro_dim KeyError) when site probs come from chained indexing T[a][b]

Parallel enumeration: silently wrong posterior weights (or a _pyro_dim KeyError) when site probs come from chained indexing T[a][b]

Description

When a discrete parent is enumerated (infer={"enumerate": "parallel"}), the parent value carries an enumeration dimension. If a downstream site's probs come from chained indexing of a table , i.e.,T[a][b], the enumeration machinery computes the wrong posterior weights, silently. The equivalent single-op indexing T[a, b] gives the correct posterior. With chained indexing on two enumerated parents the model crashes outright (KeyError: '_pyro_dim_N', funsor internal dim).

The wrong result is especially nasty because nothing raises, you get draws that look normal but come from the wrong posterior.

Repro

import jax
import jax.numpy as jnp
import numpyro
import numpyro.distributions as dist
from numpyro.infer import Predictive

P = jnp.array([
    [[0.90, 0.10], [0.70, 0.30]],  # P(wet | rain=0, u)
    [[0.40, 0.60], [0.05, 0.95]],  # P(wet | rain=1, u)
])

def model_chained():
    rain = numpyro.sample(
        "rain",
        dist.Categorical(probs=jnp.array([0.8, 0.2])),
        infer={"enumerate": "parallel"},
    )
    numpyro.sample("wet", dist.Categorical(probs=P[rain][1]), obs=1)  # chained

def model_single():
    rain = numpyro.sample(
        "rain",
        dist.Categorical(probs=jnp.array([0.8, 0.2])),
        infer={"enumerate": "parallel"},
    )
    numpyro.sample("wet", dist.Categorical(probs=P[rain, 1]), obs=1)  # single op

for name, model in [("chained", model_chained), ("single", model_single)]:
    draws = Predictive(model, num_samples=100_000, infer_discrete=True)(
        jax.random.PRNGKey(0)
    )
    print(name, float((draws["rain"] == 1).mean()))

Expected by hand: P(wet=1 | rain=0, u=1) = 0.30, P(wet=1 | rain=1, u=1) = 0.95, so

P(rain=1 | wet=1, u=1) = 0.2·0.95 / (0.8·0.30 + 0.2·0.95) = 0.4419

Observed:

chained 0.28344   <- wrong
single  0.44020   <- correct

What I've tested (numpyro 0.21.0, jax 0.11.0, funsor 0.4.7)

Note that numpyro.ops.indexing.Vindex fixes the single-getitem form but not the chained form, and no documentation distinguishes the two forms (in plain JAX they are equivalent).

Also reproduced through NUTS enumeration (mixed model with a continuous latent, discrete parent enumerated, conditional Predictive pass to recover the discrete draws): chained 0.259 vs single 0.422 (exact 0.442).

Plain-int indexing (no enumeration, e.g. forward Predictive or external factor-based engines calling the same expression with concrete values) is unaffected, the bug is specific to the enumerated-value path.

Expected

Chained and single-op indexing should be equivalent, as they are in plain JAX. Batched ((E, K)) Categorical probs from a gather are branch-aligned either way; the enumeration machinery appears to mis-handle the two-op form.

How I ran into it

I am developing a library on top of numpyro where the user can provide a dist of callables the enumerated parent values. TABLE[p1][p2] is the natural way users write multi-parent CPTs. I found this bug while working with Deepseek V4-Pro on the feature because our tests failed.

I am not an expert in Numpyro, but let me know if I can help with anything. In my code, I have documented it and used a workaround, but this feels like a real bug that should be fixed in numpyro.

Start with the parallel enumeration path used by Predictive and the indexing behavior in numpyro.ops.indexing.Vindex; compare chained and single-op indexing using the provided reproduction. Trace how enumerated values flow through P[rain][1] and P[rain][u]. Done means both indexing forms produce equivalent posterior weights and the two-enumerated-parent case no longer raises the reported _pyro_dim KeyError.