Structure learning with combinatorial variables sampling

[pawel-czyz]

In the structure learning problems there appear combinatorial latent variables (e.g., a DAG or a phylogenetic tree) which need to be inferred.

In particular, there are various MCMC kernels designed to sample these structures (e.g., partition MCMC for DAGs or different tree kernels).
I wonder if adding them to BlackJAX would be of interest.

@grfrederic and I have been thinking about rooted trees inference (which are simpler structures than phylogenetic trees or DAGs, but are already useful in evolutionary research) and it seems possible to implement them in bare JAX (on the other hand, I do not know how to do it in PyTensor).

Do you think we should create a pull request with this or it is specific enough such that it should not be included in BlackJAX?

[albcab]

I think no inference algorithm is too specific for blackjax. Focusing on blackbox, generally applicable algorithms might even be detrimental to the development of new, efficient ways to solve specific (but interesting) inference problems. So, by all means, implement any of these methods in blackjax!

I would recommend to start by opening an issue where you propose the implementation of a specific algorithm of this kind, where we can discuss the caveats of the implementation. For example, how to handle graphs as pytrees? do you use any dependencies to do this? Even if there is nothing to discuss, it's a good way to let people know you are working on a specific algorithm. Then, create a pull request.

Also if you've applied this algorithm on an evolutionary model, a great addition would be a notebook in the sampling book where you illustrate its use once its done.