The PyMC initial point function is not thread safe

[ricardoV94]

This is causing a surprising failure that None has no shape in the raveling wrapper:

nutpie/python/nutpie/compile_pymc.py

Line 71 in 6329692

n = int(np.prod(initial_value.shape))

I suspect that threads are causing the PyTensor function to clear the output storage (which is done with None) before the
thread that called it had time to retrieve the values for itself.

I guess it could also happen to intermediate storage and fail inside a PyTensor call.

[lucianopaz]

Here's the minimal reproducible example that we managed to distill

import nutpie
import pymc as pm

coords = {
    "dim1": range(50),
    "dim2": range(50),
}
with pm.Model(coords=coords) as model:
    sigma = pm.HalfNormal("sigma", default_transform=None, dims=["dim1", "dim2"])

compiled = nutpie.compile_pymc_model(model)

idata = nutpie.sample(
    compiled,
    tune=1,
    draws=1,
    chains=10,
    seed=1234,
)

This fails truly at random. The reason this fails is that there is a race condition for calling the same underlying pytensor function to produce initial points from several threads. The call that happens in different threads is here, and the threadpool is spawned here using rayon.spawn_fifo.

[lucianopaz]

I think that there two solutions to this on the rust side of things:

1. The crudest solution would be to take the initial points for all chains sequentially outside of the scope.spawn_fifo and pass them into the threads somehow.
2. Maybe a better approach would be to include a lock, just how nuts-rs does for writing to the trace here and here

Without having to go into rust, another solution would be to compile the make_initial_point_fn use a different backend and have nutpie wrap the jit_fn (hoping this is thread safe?!) instead of using pytensor's function, which still has a VM running and garbage collecting it. This is way too much work.

[aseyboldt]

Oh, I fixed this, but only pushed it to the transformation PR: 9a1da91 🤦
I'll make a PR with that fix, so that we can quickly release a new minor version...

[lucianopaz]

Oh, I fixed this, but only pushed it to the transformation PR: 9a1da91 🤦 I'll make a PR with that fix, so that we can quickly release a new minor version...

Great! We were also wondering if the initial point’s logp is checked anywhere? We found the nuts-rs lines that retry the call to the initial point function, but we didn’t see where it was checking if the initial point had a finite logp.

[aseyboldt]

It calls Chain.set_position:
https://github.com/pymc-devs/nuts-rs/blob/3d845ac09caeb1a2ff6b41e991e7c2e1ba7781b5/src/chain.rs#L162
And that will return an error if the logp is invalid. But the actual check happens in the backend-specific logp functions, for instance here: https://github.com/pymc-devs/nutpie/blob/main/src/pyfunc.rs#L157

[ricardoV94]

It calls Chain.set_position: https://github.com/pymc-devs/nuts-rs/blob/3d845ac09caeb1a2ff6b41e991e7c2e1ba7781b5/src/chain.rs#L162 And that will return an error if the logp is invalid. But the actual check happens in the backend-specific logp functions, for instance here: https://github.com/pymc-devs/nutpie/blob/main/src/pyfunc.rs#L157

May be worth a comment inline ;)