Dynamic shaping, "round" function, JAX, and a "few" more questions

Hi, I’m trying to reproduce this example (LKJ Cholesky Covariance Priors for Multivariate Normal Models — PyMC example gallery) for one of my projects. When, to test the code, I pass the dictionary returned by model.initial_point() as the initvals parameter of the pm.sample sampler, I get the following error:

Auto-assigning NUTS sampler...
Initializing NUTS using jitter+adapt_diag...
Traceback (most recent call last):

  File ~/InstProg/anaconda3/envs/myenv/lib/python3.10/site-packages/spyder_kernels/py3compat.py:356 in compat_exec
    exec(code, globals, locals)

  File ~/Bureau/detector_work/IPC/pymc_ipc.py:304
    idata = pm.sample(draws=1000, tune=1000, chains=2, cores=2,

  File ~/InstProg/anaconda3/envs/myenv/lib/python3.10/site-packages/pymc/sampling/mcmc.py:682 in sample
    initial_points, step = init_nuts(

  File ~/InstProg/anaconda3/envs/myenv/lib/python3.10/site-packages/pymc/sampling/mcmc.py:1327 in init_nuts
    initial_points = _init_jitter(

  File ~/InstProg/anaconda3/envs/myenv/lib/python3.10/site-packages/pymc/sampling/mcmc.py:1204 in _init_jitter
    ipfns = make_initial_point_fns_per_chain(

  File ~/InstProg/anaconda3/envs/myenv/lib/python3.10/site-packages/pymc/initial_point.py:86 in make_initial_point_fns_per_chain
    make_initial_point_fn(

  File ~/InstProg/anaconda3/envs/myenv/lib/python3.10/site-packages/pymc/initial_point.py:134 in make_initial_point_fn
    sdict_overrides = convert_str_to_rv_dict(model, overrides or {})

  File ~/InstProg/anaconda3/envs/myenv/lib/python3.10/site-packages/pymc/initial_point.py:47 in convert_str_to_rv_dict
    initvals[rv] = model.rvs_to_transforms[rv].backward(initval, *rv.owner.inputs)

  File ~/InstProg/anaconda3/envs/myenv/lib/python3.10/site-packages/pymc/distributions/transforms.py:157 in backward
    return pt.set_subtensor(value[..., self.diag_idxs], pt.exp(value[..., self.diag_idxs]))

IndexError: only integers, slices (`:`), ellipsis (`...`), numpy.newaxis (`None`) and integer or boolean arrays are valid indices

It doesn’t happen if I pass “chol” directly to initvals and not the transformed parameter “chol_cholesky-cov-packed__”. Are you able to reproduce the same error on your side, and could you please explain where it comes from?

Thank you!

May be better to open a separate discord thread for this one, and try to include the smallest reproducible example you can come up with.

Thank you @ricardoV94 for your reply. I have just created a separate thread about this error I got. (Sampler initialization error with a model containing an LKJCholeskyCov distribution)

Hi @ricardoV94,

We now have a PyMC model that works well, especially since you fixed that nested CustomDist error. Thank you again! We have made good progress testing our model using simulated data and plan to apply it to real data generated from our experiment/detector very soon. It is an exciting time for us!

I would still be please interested to get your help/comments about questions Q4 and Q5 from my previous message in order to improve and get the best out of my current PyMC model, and also to a lesser extent about questions Q1 and Q2 to make sure I understand how PyTensor works and that I am using it properly.

Thank you again for your help and time.

Yes, this helper remove_value_transforms will return an equivalent model without any transforms, whose logp functions will be defined on untransformed variables: pymc.model.transform.conditioning — PyMC 5.10.0 documentation

2 Likes

Many thanks, @ricardoV94!

I have just tested it, it works great!
I was also wondering please:
1-how do we feed the initvals parameter to the sampler in case we want to keep the parameters transformed?
2- how to apply the logp on new “observed” data - by providing the observed data as MutableData and using pm.set_data?
Thank you!

I think you can pass either transformed or untransformed initvals to pm.sample and the sampler will use those. Not sure right now.

That’s probably the simplest approach yes

1 Like

Thank you for your quick reply!