Forward model argument to add_fit_to_inference_data in find_MAP

[jessegrabowski]

Bugfix for find_MAP. The following code currently errors:

    with pm.Model() as m:
        mu = pm.Normal("mu", 0, 1)
        sigma = pm.Exponential("sigma", 1)
        y_hat = pm.Normal("y_hat", mu=mu, sigma=sigma, observed=np.random.normal(size=10))

    idata = find_MAP(
        model=m,
        method="L-BFGS-B",
        use_grad=True,
        use_hess=False,
        use_hessp=False,
        progressbar=False,
        compile_kwargs={"mode": "JAX"},
    )

Because the model argument is not correctly forwarded everywhere it is needed inside find_MAP. This patches the bug and adds a regression test.

[ricardoV94]

Does this superseed: #538 ?

[ricardoV94]

How comes PyPI CI is failing with import?

[jessegrabowski]

Does this superseed: #538 ?

Yes. I didn't check if anything else was open before I did this. Sorry @OriolAbril

How comes PyPI CI is failing with import?

I think this will be fixed by #541 ?

[OriolAbril]

Can we also add the correct pin for pytensor? The map and laplace functionality needs newer pytensor than what is currently allowed. Don't worry about my PR

[jessegrabowski]

It's actually not map and laplace needing the new pytensor, it's the INLA function.

I think we could move that specific import inside this function? It shouldn't effect 99% of users.

[tomicapretto]

@jessegrabowski this is the dependency problem I encountered before

conda create --name test-pymc-extras python=3.12
conda activate test-pymc-extras
conda install -c conda-forge pymc-extras
python

import pymc_extras

Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/home/tomas/miniconda3/envs/test-pymc-extras/lib/python3.12/site-packages/pymc_extras/__init__.py", line 20, in <module>
    from pymc_extras.inference import find_MAP, fit, fit_laplace, fit_pathfinder
  File "/home/tomas/miniconda3/envs/test-pymc-extras/lib/python3.12/site-packages/pymc_extras/inference/__init__.py", line 17, in <module>
    from pymc_extras.inference.laplace_approx.laplace import fit_laplace
  File "/home/tomas/miniconda3/envs/test-pymc-extras/lib/python3.12/site-packages/pymc_extras/inference/laplace_approx/laplace.py", line 38, in <module>
    from pytensor.tensor.optimize import minimize
ModuleNotFoundError: No module named 'pytensor.tensor.optimize'

pymc.__version__, pytensor.__version__
# ('5.23.0', '2.31.3')

and the pymc_extras I got installed was 0.3.1

If I do

conda install -c conda-forge "pytensor=2.31.7"
python

import pymc_extras
pymc_extras.__version__
# '0.3.1'

I think the correct pin for pytensor is >=2.31.4 (and some upper pin)

[jessegrabowski]

Thanks Tomi. I updated the pins. I didn't add an upper pin -- I don't really like them, and @maresb said he doesn't either, so that's confirmation bias :D

[jessegrabowski]

I moved the version pin and dependency stuff to #542 , they didn't seem appropriate in this PR

[codecov-commenter]

Codecov Report

Attention: Patch coverage is 0% with 2 lines in your changes missing coverage. Please review.

Project coverage is 64.15%. Comparing base (9c50e94) to head (1d75f75).

Files with missing lines	Patch %	Lines
pymc_extras/inference/laplace_approx/find_map.py	0.00%	2 Missing ⚠️

Additional details and impacted files

@@             Coverage Diff             @@
##             main     #543       +/-   ##
===========================================
+ Coverage   51.83%   64.15%   +12.32%     
===========================================
  Files          55       65       +10     
  Lines        6048     6487      +439     
===========================================
+ Hits         3135     4162     +1027     
+ Misses       2913     2325      -588     

Files with missing lines	Coverage Δ
pymc_extras/inference/laplace_approx/find_map.py	20.58% <0.00%> (ø)

... and 31 files with indirect coverage changes

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[tomicapretto]

I moved the version pin and dependency stuff to #542 , they didn't seem appropriate in this PR

Makes sense!

[tomicapretto]

Is it relevant / worth checking the presence of any data group? I see above a call like the following within find_MAP

    idata = add_data_to_inference_data(
        idata=idata, progressbar=False, model=model, compile_kwargs=compile_kwargs
    )

[jessegrabowski]

crap yes it is, but I just clicked merge

[tomicapretto]

haha, well another thing I just saw:

Is it necessary to compute deterministics within add_data_to_inference_data? See

pymc-extras/pymc_extras/inference/laplace_approx/idata.py

Lines 211 to 224 in fd2933f

	if model.deterministics:
	expand_dims = {}
	if "chain" not in idata.posterior.coords:
	expand_dims["chain"] = [0]
	if "draw" not in idata.posterior.coords:
	expand_dims["draw"] = [0]
	idata.posterior = pm.compute_deterministics(
	idata.posterior.expand_dims(expand_dims),
	model=model,
	merge_dataset=True,
	progressbar=progressbar,
	compile_kwargs=compile_kwargs,
	)

I guess that function is not exposed to the user, but I just wanted to raise that potentially silent side effect.

[jessegrabowski]

it's not necessary, but that's what pm.sample does right?

[tomicapretto]

Ahh I see now why it's included