Very interesting, and we would love to help, but to better advise maybe you could provide a bit of pseudocode as an example?
A short answer for now:
- PyMC4 is no more and moving forward everything will be in GitHub - pymc-devs/pymc: Probabilistic Programming in Python: Bayesian Modeling and Probabilistic Machine Learning with Aesara (aka PyMC v4 will be the next major release).
- For that I would like to understand a bit more of how the algorithm works, from MuseInference.jl · MuseInference it seems it requires prior log_prob for the \theta, joint log likelihood and its gradient to get the Score Marginal likelihood function, and a prior simulation function to get the approximation for the score function - I think the closest thing might be the implementation for Sequential Monte Carlo which isolate similar component from a PyMC model: pymc/pymc/smc at main · pymc-devs/pymc · GitHub