PyMC Implementation of Pathfinder VI

[aphc14]

Summary:

• Adds a PyMC implementation of Pathfinder Variational Inference using PyTensor operations. The new implementation allows users to choose between PyMC and BlackJAX backends while maintaining the same API.
• Added lbfgs.py module implementing L-BFGS optimisation with history tracking
• Extended pathfinder.py with PyMC implementation using PyTensor operations.

Note: Another draft PR will be sent that focuses on a PyTensor symbolic implementation using pytensor.function. I've sent two PR drafts to get feedback on which version would be better.

import pymc as pm
import pymc_experimental as pmx

with pm.Model() as model:
    # ... model definition ...
    idata = pmx.fit(
        method="pathfinder",
        inference_backend="pymc",  # or "blackjax"
        random_seed=42
    )

with model:
    # eight_schools_model
    idata = pmx.fit(model=model, method="pathfinder", random_seed=41, inference_backend="pymc")
    
    # New implementation now passes this assertion! :)
    # FIXME: pathfinder doesn't find a reasonable mean! Fix bug or choose model pathfinder can handle
    np.testing.assert_allclose(idata.posterior["mu"].mean(), 5.0, atol=1.0)

    # But, it also fails this :(
    # FIXME: now the tau is being underestimated. getting tau around 1.5.
    # np.testing.assert_allclose(idata.posterior["tau"].mean(), 4.15, atol=0.5)

[twiecki]

This looks great @aphc14, can you find me on linkedin?

[aphc14]

@twiecki yup, added you :)
@fonnesbeck, FYI, Multipath Pathfinder has just been implemented!

[aphc14]

closed this in favour of #387