How can we calculate the log partition function, given that we can obtain node marginals through inferer.get_beliefs? Is it also feasible to acquire factor marginals for this purpose? While the typical approach to compute the log partition function is as follows:

$$log Z \approx \sum_{i\in V} (1-d_i)H_i + \sum_{(i,j)\in E} I_{ij}$$

This formula requires both the entropy $H_i$ for individual nodes and the mutual information $I_{ij}$ for edges. To compute these, we need the node marginals $b_i(x_si)$ and the joint marginal probabilities $b_{ij}(x_i,x_j)$.

I just attempted to install PGMax via pip: i.e. pip install pgmax. While it did work, I noticed that it didn't seem to install any of the dependencies. Indeed, when I opened up a terminal and ran from pgmax import fgraph, I got a ModuleNotFoundError: No module named 'jax'.

The core issue seems to be that there are no requirements listed in the setup.py file, and the requirements.txt file is also not linked to setup.py in any way (indeed, this commit seems to have removed the link between setup.py and requirements.txt). If I clone the repo and run pip install -r requirements.txt, then things work fine, but the pip installation does still seem currently broken.

Hi.
When running notebooks in colab, for example rcn.ipynb, there is a problem in the installation cell:

This is because !cd PGMax works locally and not globally.
Solution: replace this line to %cd PGMax
P.S: I can make PR for this.

• In examples/gmrf.ipynb, maybe replace from jax.example_libraries import optimizers with optax. Also the term "GMRF" usually refers to Gaussian MRF, not grid :) Finally the learning code is super slow, even on an A100... (public colab)
• in examples/rbm.ipynb, you mix jax.vmap with np.random, which seems dangerous...
• in examples/ising_model.ipynb, the line batch_loss = jax.jit(jax.vmap(loss, in_axes=(None, {variables: 0}), out_axes=0)) is very obscure. What is the dict with a key called 'variables'? And in what sense are these evidence / potential 'updates', as opposed to just static things?
• it would be helpful to add a link to the ising example from https://github.com/deepmind/PGMax/tree/main#getting-started

I'm trying to use named variables using vgroup.Vardict, but without much success. I cannot find any reference to or usage of Vardict. Specifically I'm running into an issue when trying to update evidence.

In pgmax/vgroup/vdict.py the function flatten expects a mapping from the varnames to e.g. evidence, but in pgmax/infer/bp_state.py only the evidence is passed -> flat_data = name.flatten(data)

An example snippet:

root_variables = vgroup.VarDict(num_states=2, variable_names=("a"))
leaf_variables = vgroup.VarDict(num_states=2, variable_names=("b"))
fg = fgraph.FactorGraph(variable_groups=[root_variables, leaf_variables])

r1_factor = factor.EnumFactor(
      variables=[root_variables["a"]],
      factor_configs=np.array([[0], [1]]),
      log_potentials=np.array([math.log(0.9), math.log(0.1)]),
)
fg.add_factors(r1_factor)

pairwise_factor = factor.EnumFactor(
    variables=[root_variables["a"],leaf_variables["b"]],
    factor_configs=np.array(list(itertools.product(np.arange(2), repeat=2))),
    log_potentials=np.array([math.log(0.9), math.log(0.1), math.log(0.2), math.log(0.8)]),
)
fg.add_factors(pairwise_factor)

evidence_updates={root_variables: np.array([math.log(0.9), math.log(0.1)])}
inferer = infer.build_inferer(fg.bp_state, backend="bp")
inferer_arrays = inferer.init(evidence_updates=evidence_updates)

Snippet fails at the init step when trying to update evidence.
I'm not sure if this is an actual issue or whether I'm simply wrong in the way I created the graph.

It would be great to have an example of inference in a directed graphical model. e.g. just the standard sprinkler example would be really useful for getting a handle on how to use the package.

Dear PGMax Team,

first of all: Thank you very much for the library, it works super well and we found it to be orders of magnitude faster than other libraries. Furthermore, the flattened potentials are super neat :)

If you don't mind, I would have two questions and would be very grateful if you could provide some pointers:

• prob == 0 we have unary potentials that can be (numerically) zero. This leads to errors where the prediction seems to completely break down for some variables, although there are variable -> state assignments that would not hit the prob == 0. Initially, we replaced the log potentials at -np.inf with your pgmax.utils.NEG_INF = -1e20. While this fixed some issues, we still got random assignments in some cases that seem to disappear if we use say -20 as the smallest unary potential. I have a hunch that this happens specifically for temperature=0 when getting map_states but couldn't investigate. What is the guidance on this? Could you point towards a good paper that has some experiments regarding this? Is there an intuitive explanation why LBP fails in this case?

• Initialization. Currently, we use bp_arrays = bp.run_bp(bp.init(), num_iters=num_iters, damping=damping) as given in most tutorials. This seems to init the log_potentials as 0, right? Is there some way we could use prior knowledge to do this? Maybe init with the most likely states given the unary potentials?

• If you have any other papers that give good hints on how to squeeze the most out of LBP, I would be happy to read them.

Best and Thanks
Lukas