Open nathanielvirgo opened 2 years ago
This is definitely another part of the documentation we need to improve.
You can update evidence using intuitive assignment syntax. https://github.com/vicariousinc/PGMax/blob/f7f3e3fe7fbb3a3c7582b7093c6ecd56eb04963a/examples/ising_model.py#L96 Here is an example for an Ising model. You index the evidence with variable names (tuples of the type (group_name, variable_name_in_group)
) and do assignment. If you are working with variable groups (e.g. an NDVariableArray
), you can also use intuitive assignment syntax with arrays (e.g. https://github.com/vicariousinc/PGMax/blob/f7f3e3fe7fbb3a3c7582b7093c6ecd56eb04963a/examples/ising_model.py#L103).
The evidence_updates
argument to run_bp
updates evidence on the fly for the given belief propagation run. You can use the elements in this dictionary in the above assignment syntax to achieve the same kind of evidence updates.
Some additional comments for this two different ways of doing evidence updates:
run_bp
later you will always work with the updated evidence.evidence_updates
the updates are not persistent. It only updates the evidence for that particular run_bp
run. But by structuring evidence_updates
as a dictionary of jnp arrays and having run_bp
as pure python functions, we obtain the ability to vectorize over/get gradients w.r.t. the evidence updates using jax.vmap
/jax.grad
. This is the main motivation for organizing things this way.Also you are more than welcome to contribute an example with directed models (since you obviously are working on one). Please refer to the developer installation instructions and contribution guide, and feel free to open a PR when it's ready! I can make more comments on the PR and make it part of our examples.
I would like to do the following, which is one step of the EM algorithm:
I'm currently stuck at the second bullet point. If I search the docs for "evidence" I find
pgmax.fg.graph.Evidence
andpgmax.fg.graph.update_evidence
, but the documentation for these is very incomplete and doesn't specify the format in which the inputs should be provided or the procedure for using them.The blog post doesn't use these but something else - it uses an argument
evidence_updates
torun_bp
, but this is not mentioned in the documentation. I can't really understand what this argument is doing in the example, but I suspect it's not what I'm looking for.So I'm a bit stuck on understanding the basic usage, and I'm wondering if I could either get a hint about it, and/or whether the documentation could be updated to make it a bit clearer.