use a swich condition for belief propagation

[LegrandNico]

The belief propagation function consists of a for loop, which might be sub-optimal when caching the function. It could instead be coded as a switch node where relevant update functions are indexed.

[LegrandNico]

Because jax.lax.switch turns everything into Tracers, the node_idx cannot be set as static anymore and this would require using Partial beforehand. I am not sure there is any benefit to doing this in the current codebase, but this is worth trying in the vectorized version, where node_idx can be a tracer (see #124 ).