Sample output object with Nutpie

Questions v5
,
1

Hi,

(Im using PyMC ver 5.19 and Windows 10)

Im using nutpie to compile and sample a model (previously run fine with pm.sample etc). This works very well (much quicker) but the resulting trace object is different (previously I had pm.sample return an InferenceData object).

Is there a straightforward way to create the exact same output object with nutpie.sample, i.e. include the extra attributes (e.g. ‘constant_data’ & ‘observed_data’) as in:

Inference data with groups:
> posterior
> sample_stats
> observed_data
> constant_data

(Ive tried doing this by hand but its painful).

Thanks

2

You can perhaps use these helpers? pymc/pymc/sampling/jax.py at 82716fbee6d722e4552f88f998e1a20f84045c47 · pymc-devs/pymc · GitHub

3

Works perfectly, thank you