Something wrong in the code from `PyMC 4.0 with labeled coords and dims`

Questions
3

For the record, the devs did pick a computational backend – it’s pytensor only moving forward. That said, there is a lot of old content out there that was made using theano and aesara. Luckily, it’s all mostly just drop-in replacements.

This particular issue seems to be a new bug related to indexing the same child random variable multiple ways in the same model. It’s quite niche, because you have to 1) define a random variable, 2) perform some intermediate computations on it, then 3) index it in multiple different ways. I ran into it myself yesterday, it seems to be quite fresh.

EDIT: I was unable to reproduce it on newer versions, see below. My testing was initially done on Mac OS running PyMC version 5.0.1

As an interim solution, you can avoid the error by compiling the model to JAX or Rust. JAX is only an option if you’re not on windows, but if you aren’t you can add:

import pymc.sampling

with pm.Model(coords=coords) as rugby_model:
    # nothing changes in model definition

    rugby_idata = pymc.sampling_jax.sample_numpyro_nuts()

If you are on windows, you can use nutpie to sample with rust. Remove the pm.sample line entirely, then outside the model context compile to rust and sample:

with pm.Model(coords=coords) as rugby_model:
    # nothing changes in model definition
    # don't run pm.sample!

import nutpie
numba_mod = nutpie.compile_pymc_model(rugby_model)
rugby_idata = nutpie.sample(numba_mod)

I believe the API for accessing these alternative samplers has improved or will be improved soon, but I’m still on 5.0.1. These are all just work-arounds, because you did stumble on a bug. But they’re also nice to know, because they are often much faster than the C backend!

EDIT: After some more tests, I could not reproduce the bug on my Windows computer running pymc 5.0.2. What version/platform are you running?

1 Like