Use of dict.pop in pm.sample leads to side effects

[jessegrabowski]

Description

I commonly create a sample_kwargs variable in notebooks and re-use it in multiple models. Such a dictionary ends up being modified in-place by pm.sample when entries are popped, for example here and here, leading to models beyond the first silently using defaults settings.

I guess the easiest fix would just be to deepcopy kwargs, so that there are no side effects on user inputs. Using pop here is a bit overkill (why not just ,get?) but deepcopy is a 1-line change.

[inclinedadarsh]

Hey, this seems to be like a good place to start contributing. I'd like to attempt it.

As much as I have understood, I'll have to deepcopy the kwargs right before the files you have linked.

Can you please help me by redirecting me to sample/example notebooks using which I can reproduce this issue?

[jessegrabowski]

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, shape=(1000,))
    
data = pm.draw(y_hat)
m = pm.observe(m, {'y_hat':data})

nuts_sampler_kwargs = {'backend':'jax', 'gradient_backend':'jax'}
sample_kwargs = {'nuts_sampler':'nutpie', 'nuts_sampler_kwargs':nuts_sampler_kwargs}

with m:
    assert sample_kwargs['nuts_sampler_kwargs'] == {'backend':'jax', 'gradient_backend':'jax'}
    idata = pm.sample(**sample_kwargs)
    assert sample_kwargs['nuts_sampler_kwargs'] == {'backend':'jax', 'gradient_backend':'jax'}

The 2nd assert should pass