Error messages when using the pymc or nutpie NUTS samplers in combination with pymc-experimental

[rklees]

When I use the pymc or nutpie NUTS samplers in combination with pymc-experimental I get error messages which I can't solve. An example is shown below. Is there anyone who knows how to fix this?

---

RemoteTraceback Traceback (most recent call last)
RemoteTraceback:
"""
Traceback (most recent call last):
File "/Library/Frameworks/Python.framework/Versions/3.11/lib/python3.11/site-packages/pytensor/compile/function/types.py", line 970, in call
self.vm()
File "/Library/Frameworks/Python.framework/Versions/3.11/lib/python3.11/site-packages/pytensor/scan/op.py", line 1648, in rval
r = p(n, [x[0] for x in i], o)
^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/Library/Frameworks/Python.framework/Versions/3.11/lib/python3.11/site-packages/pytensor/scan/op.py", line 1576, in p
t_fn, n_steps = scan_perform_ext.perform(
^^^^^^^^^^^^^^^^^^^^^^^^^
File "pytensor/scan/scan_perform.pyx", line 397, in pytensor.scan.scan_perform.perform
AttributeError: 'ArrayImpl' object has no attribute 'data'

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
File "/Library/Frameworks/Python.framework/Versions/3.11/lib/python3.11/site-packages/pymc/sampling/parallel.py", line 122, in run
self._start_loop()
File "/Library/Frameworks/Python.framework/Versions/3.11/lib/python3.11/site-packages/pymc/sampling/parallel.py", line 174, in _start_loop
point, stats = self._step_method.step(self._point)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/Library/Frameworks/Python.framework/Versions/3.11/lib/python3.11/site-packages/pymc/step_methods/arraystep.py", line 174, in step
return super().step(point)
^^^^^^^^^^^^^^^^^^^
File "/Library/Frameworks/Python.framework/Versions/3.11/lib/python3.11/site-packages/pymc/step_methods/arraystep.py", line 100, in step
apoint, stats = self.astep(q)
^^^^^^^^^^^^^
File "/Library/Frameworks/Python.framework/Versions/3.11/lib/python3.11/site-packages/pymc/step_methods/hmc/base_hmc.py", line 168, in astep
start = self.integrator.compute_state(q0, p0)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/Library/Frameworks/Python.framework/Versions/3.11/lib/python3.11/site-packages/pymc/step_methods/hmc/integration.py", line 56, in compute_state
logp, dlogp = self._logp_dlogp_func(q)
^^^^^^^^^^^^^^^^^^^^^^^^
File "/Library/Frameworks/Python.framework/Versions/3.11/lib/python3.11/site-packages/pymc/model/core.py", line 378, in call
cost, *grads = self._pytensor_function(*grad_vars)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/Library/Frameworks/Python.framework/Versions/3.11/lib/python3.11/site-packages/pytensor/compile/function/types.py", line 983, in call
raise_with_op(
File "/Library/Frameworks/Python.framework/Versions/3.11/lib/python3.11/site-packages/pytensor/link/utils.py", line 531, in raise_with_op
raise exc_value.with_traceback(exc_trace)
File "/Library/Frameworks/Python.framework/Versions/3.11/lib/python3.11/site-packages/pytensor/compile/function/types.py", line 970, in call
self.vm()
File "/Library/Frameworks/Python.framework/Versions/3.11/lib/python3.11/site-packages/pytensor/scan/op.py", line 1648, in rval
r = p(n, [x[0] for x in i], o)
^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/Library/Frameworks/Python.framework/Versions/3.11/lib/python3.11/site-packages/pytensor/scan/op.py", line 1576, in p
t_fn, n_steps = scan_perform_ext.perform(
^^^^^^^^^^^^^^^^^^^^^^^^^
File "pytensor/scan/scan_perform.pyx", line 397, in pytensor.scan.scan_perform.perform
AttributeError: 'ArrayImpl' object has no attribute 'data'
Apply node that caused the error: Scan{grad_of_forward_kalman_pass, while_loop=False, inplace=none}(144, Transpose{axes=[0, 2, 1]}.0, Transpose{axes=[0, 2, 1]}.0, Transpose{axes=[0, 2, 1]}.0, ExpandDims{axis=1}.0, ExpandDims{axis=1}.0, [[-7.18797 ... 9693e+01]], Subtensor{start:stop:step}.0, Subtensor{start:stop:step}.0, Subtensor{start:stop:step}.0, Subtensor{start:stop:step}.0, [[0. 0. 0. ... 0. 0. 0.]], [[[0. 0. 0 ... . 0. 0.]]], [[[0. 0. 0 ... . 0. 0.]]], 144, 144, transition, DropDims{axis=0}.0, Transpose{axes=[1, 0]}.0, Dot22.0, state_cov.T, selection, selection.T)
Toposort index: 94
Inputs types: [TensorType(int64, shape=()), TensorType(float64, shape=(None, 6, 6)), TensorType(float64, shape=(None, None, None)), TensorType(float64, shape=(None, 6, 6)), TensorType(float64, shape=(None, 1, 6)), TensorType(float64, shape=(None, 1, 6)), TensorType(float64, shape=(144, 1)), TensorType(float64, shape=(None, 6, 6)), TensorType(float64, shape=(None, 1)), TensorType(float64, shape=(None, 6, 6)), TensorType(float64, shape=(None, None, None)), TensorType(float64, shape=(145, 6)), TensorType(float64, shape=(145, 6, 6)), TensorType(float64, shape=(2, 5, 5)), TensorType(int64, shape=()), TensorType(int64, shape=()), TensorType(float64, shape=(6, 6)), TensorType(float64, shape=(1, 1)), TensorType(float64, shape=(6, 6)), TensorType(float64, shape=(5, 6)), TensorType(float64, shape=(5, 5)), TensorType(float64, shape=(6, 5)), TensorType(float64, shape=(5, 6))]
Inputs shapes: [(), (144, 6, 6), (144, 1, 1), (144, 6, 6), (144, 1, 6), (144, 1, 6), (144, 1), (144, 6, 6), (144, 1), (144, 6, 6), (144, 1, 1), (145, 6), (145, 6, 6), (2, 5, 5), (), (), (6, 6), (1, 1), (6, 6), (5, 6), (5, 5), (6, 5), (5, 6)]
Inputs strides: [(), (-288, 8, 48), (-8, 8, 8), (-288, 8, 48), (-48, 48, 8), (-48, 48, 8), (-8, 8), (-288, 48, 8), (-8, 8), (-288, 48, 8), (-8, 8, 8), (48, 8), (288, 48, 8), (200, 40, 8), (), (), (48, 8), (8, 8), (8, 48), (48, 8), (8, 40), (40, 8), (8, 40)]
Inputs values: [array(144), 'not shown', 'not shown', 'not shown', 'not shown', 'not shown', 'not shown', 'not shown', 'not shown', 'not shown', 'not shown', 'not shown', 'not shown', 'not shown', array(144), array(144), 'not shown', array([[0.02588485]]), 'not shown', 'not shown', 'not shown', 'not shown', 'not shown']
Outputs clients: [[Subtensor{start:stop:step}(Scan{grad_of_forward_kalman_pass, while_loop=False, inplace=none}.0, 144, 143, -1)], [Subtensor{start:stop:step}(Scan{grad_of_forward_kalman_pass, while_loop=False, inplace=none}.1, 144, 143, -1)], [Subtensor{i}(Scan{grad_of_forward_kalman_pass, while_loop=False, inplace=none}.2, -1)], [Reshape{1}(Scan{grad_of_forward_kalman_pass, while_loop=False, inplace=none}.3, [-1])], [Reshape{1}(Scan{grad_of_forward_kalman_pass, while_loop=False, inplace=none}.4, [144])]]

HINT: Re-running with most PyTensor optimizations disabled could provide a back-trace showing when this node was created. This can be done by setting the PyTensor flag 'optimizer=fast_compile'. If that does not work, PyTensor optimizations can be disabled with 'optimizer=None'.
HINT: Use the PyTensor flag exception_verbosity=high for a debug print-out and storage map footprint of this Apply node.
"""

The above exception was the direct cause of the following exception:

AttributeError Traceback (most recent call last)
AttributeError: 'ArrayImpl' object has no attribute 'data'
Apply node that caused the error: Scan{grad_of_forward_kalman_pass, while_loop=False, inplace=none}(144, Transpose{axes=[0, 2, 1]}.0, Transpose{axes=[0, 2, 1]}.0, Transpose{axes=[0, 2, 1]}.0, ExpandDims{axis=1}.0, ExpandDims{axis=1}.0, [[-7.18797 ... 9693e+01]], Subtensor{start:stop:step}.0, Subtensor{start:stop:step}.0, Subtensor{start:stop:step}.0, Subtensor{start:stop:step}.0, [[0. 0. 0. ... 0. 0. 0.]], [[[0. 0. 0 ... . 0. 0.]]], [[[0. 0. 0 ... . 0. 0.]]], 144, 144, transition, DropDims{axis=0}.0, Transpose{axes=[1, 0]}.0, Dot22.0, state_cov.T, selection, selection.T)
Toposort index: 94
Inputs types: [TensorType(int64, shape=()), TensorType(float64, shape=(None, 6, 6)), TensorType(float64, shape=(None, None, None)), TensorType(float64, shape=(None, 6, 6)), TensorType(float64, shape=(None, 1, 6)), TensorType(float64, shape=(None, 1, 6)), TensorType(float64, shape=(144, 1)), TensorType(float64, shape=(None, 6, 6)), TensorType(float64, shape=(None, 1)), TensorType(float64, shape=(None, 6, 6)), TensorType(float64, shape=(None, None, None)), TensorType(float64, shape=(145, 6)), TensorType(float64, shape=(145, 6, 6)), TensorType(float64, shape=(2, 5, 5)), TensorType(int64, shape=()), TensorType(int64, shape=()), TensorType(float64, shape=(6, 6)), TensorType(float64, shape=(1, 1)), TensorType(float64, shape=(6, 6)), TensorType(float64, shape=(5, 6)), TensorType(float64, shape=(5, 5)), TensorType(float64, shape=(6, 5)), TensorType(float64, shape=(5, 6))]
Inputs shapes: [(), (144, 6, 6), (144, 1, 1), (144, 6, 6), (144, 1, 6), (144, 1, 6), (144, 1), (144, 6, 6), (144, 1), (144, 6, 6), (144, 1, 1), (145, 6), (145, 6, 6), (2, 5, 5), (), (), (6, 6), (1, 1), (6, 6), (5, 6), (5, 5), (6, 5), (5, 6)]
Inputs strides: [(), (-288, 8, 48), (-8, 8, 8), (-288, 8, 48), (-48, 48, 8), (-48, 48, 8), (-8, 8), (-288, 48, 8), (-8, 8), (-288, 48, 8), (-8, 8, 8), (48, 8), (288, 48, 8), (200, 40, 8), (), (), (48, 8), (8, 8), (8, 48), (48, 8), (8, 40), (40, 8), (8, 40)]
Inputs values: [array(144), 'not shown', 'not shown', 'not shown', 'not shown', 'not shown', 'not shown', 'not shown', 'not shown', 'not shown', 'not shown', 'not shown', 'not shown', 'not shown', array(144), array(144), 'not shown', array([[0.02588485]]), 'not shown', 'not shown', 'not shown', 'not shown', 'not shown']
Outputs clients: [[Subtensor{start:stop:step}(Scan{grad_of_forward_kalman_pass, while_loop=False, inplace=none}.0, 144, 143, -1)], [Subtensor{start:stop:step}(Scan{grad_of_forward_kalman_pass, while_loop=False, inplace=none}.1, 144, 143, -1)], [Subtensor{i}(Scan{grad_of_forward_kalman_pass, while_loop=False, inplace=none}.2, -1)], [Reshape{1}(Scan{grad_of_forward_kalman_pass, while_loop=False, inplace=none}.3, [-1])], [Reshape{1}(Scan{grad_of_forward_kalman_pass, while_loop=False, inplace=none}.4, [144])]]

HINT: Re-running with most PyTensor optimizations disabled could provide a back-trace showing when this node was created. This can be done by setting the PyTensor flag 'optimizer=fast_compile'. If that does not work, PyTensor optimizations can be disabled with 'optimizer=None'.
HINT: Use the PyTensor flag exception_verbosity=high for a debug print-out and storage map footprint of this Apply node.

The above exception was the direct cause of the following exception:

ParallelSamplingError Traceback (most recent call last)
Cell In[17], line 3
1 sampler = ["pymc", "numpyro", "nutpie", "blackjax"]
2 with model_1:
----> 3 idata = pm.sample(nuts_sampler=sampler[0], tune=1000, draws=2000, chains=4)

File /Library/Frameworks/Python.framework/Versions/3.11/lib/python3.11/site-packages/pymc/sampling/mcmc.py:802, in sample(draws, tune, chains, cores, random_seed, progressbar, step, nuts_sampler, initvals, init, jitter_max_retries, n_init, trace, discard_tuned_samples, compute_convergence_checks, keep_warning_stat, return_inferencedata, idata_kwargs, nuts_sampler_kwargs, callback, mp_ctx, model, **kwargs)
800 _print_step_hierarchy(step)
801 try:
--> 802 _mp_sample(**sample_args, **parallel_args)
803 except pickle.PickleError:
804 _log.warning("Could not pickle model, sampling singlethreaded.")

File /Library/Frameworks/Python.framework/Versions/3.11/lib/python3.11/site-packages/pymc/sampling/mcmc.py:1191, in _mp_sample(draws, tune, step, chains, cores, random_seed, start, progressbar, traces, model, callback, mp_ctx, **kwargs)
1189 try:
1190 with sampler:
-> 1191 for draw in sampler:
1192 strace = traces[draw.chain]
1193 strace.record(draw.point, draw.stats)

File /Library/Frameworks/Python.framework/Versions/3.11/lib/python3.11/site-packages/pymc/sampling/parallel.py:448, in ParallelSampler.iter(self)
445 self._progress.update(self._total_draws)
447 while self._active:
--> 448 draw = ProcessAdapter.recv_draw(self._active)
449 proc, is_last, draw, tuning, stats = draw
450 self._total_draws += 1

File /Library/Frameworks/Python.framework/Versions/3.11/lib/python3.11/site-packages/pymc/sampling/parallel.py:330, in ProcessAdapter.recv_draw(processes, timeout)
328 else:
329 error = RuntimeError(f"Chain {proc.chain} failed.")
--> 330 raise error from old_error
331 elif msg[0] == "writing_done":
332 proc._readable = True

ParallelSamplingError: Chain 1 failed with: 'ArrayImpl' object has no attribute 'data'
Apply node that caused the error: Scan{grad_of_forward_kalman_pass, while_loop=False, inplace=none}(144, Transpose{axes=[0, 2, 1]}.0, Transpose{axes=[0, 2, 1]}.0, Transpose{axes=[0, 2, 1]}.0, ExpandDims{axis=1}.0, ExpandDims{axis=1}.0, [[-7.18797 ... 9693e+01]], Subtensor{start:stop:step}.0, Subtensor{start:stop:step}.0, Subtensor{start:stop:step}.0, Subtensor{start:stop:step}.0, [[0. 0. 0. ... 0. 0. 0.]], [[[0. 0. 0 ... . 0. 0.]]], [[[0. 0. 0 ... . 0. 0.]]], 144, 144, transition, DropDims{axis=0}.0, Transpose{axes=[1, 0]}.0, Dot22.0, state_cov.T, selection, selection.T)
Toposort index: 94
Inputs types: [TensorType(int64, shape=()), TensorType(float64, shape=(None, 6, 6)), TensorType(float64, shape=(None, None, None)), TensorType(float64, shape=(None, 6, 6)), TensorType(float64, shape=(None, 1, 6)), TensorType(float64, shape=(None, 1, 6)), TensorType(float64, shape=(144, 1)), TensorType(float64, shape=(None, 6, 6)), TensorType(float64, shape=(None, 1)), TensorType(float64, shape=(None, 6, 6)), TensorType(float64, shape=(None, None, None)), TensorType(float64, shape=(145, 6)), TensorType(float64, shape=(145, 6, 6)), TensorType(float64, shape=(2, 5, 5)), TensorType(int64, shape=()), TensorType(int64, shape=()), TensorType(float64, shape=(6, 6)), TensorType(float64, shape=(1, 1)), TensorType(float64, shape=(6, 6)), TensorType(float64, shape=(5, 6)), TensorType(float64, shape=(5, 5)), TensorType(float64, shape=(6, 5)), TensorType(float64, shape=(5, 6))]
Inputs shapes: [(), (144, 6, 6), (144, 1, 1), (144, 6, 6), (144, 1, 6), (144, 1, 6), (144, 1), (144, 6, 6), (144, 1), (144, 6, 6), (144, 1, 1), (145, 6), (145, 6, 6), (2, 5, 5), (), (), (6, 6), (1, 1), (6, 6), (5, 6), (5, 5), (6, 5), (5, 6)]
Inputs strides: [(), (-288, 8, 48), (-8, 8, 8), (-288, 8, 48), (-48, 48, 8), (-48, 48, 8), (-8, 8), (-288, 48, 8), (-8, 8), (-288, 48, 8), (-8, 8, 8), (48, 8), (288, 48, 8), (200, 40, 8), (), (), (48, 8), (8, 8), (8, 48), (48, 8), (8, 40), (40, 8), (8, 40)]
Inputs values: [array(144), 'not shown', 'not shown', 'not shown', 'not shown', 'not shown', 'not shown', 'not shown', 'not shown', 'not shown', 'not shown', 'not shown', 'not shown', 'not shown', array(144), array(144), 'not shown', array([[0.02588485]]), 'not shown', 'not shown', 'not shown', 'not shown', 'not shown']
Outputs clients: [[Subtensor{start:stop:step}(Scan{grad_of_forward_kalman_pass, while_loop=False, inplace=none}.0, 144, 143, -1)], [Subtensor{start:stop:step}(Scan{grad_of_forward_kalman_pass, while_loop=False, inplace=none}.1, 144, 143, -1)], [Subtensor{i}(Scan{grad_of_forward_kalman_pass, while_loop=False, inplace=none}.2, -1)], [Reshape{1}(Scan{grad_of_forward_kalman_pass, while_loop=False, inplace=none}.3, [-1])], [Reshape{1}(Scan{grad_of_forward_kalman_pass, while_loop=False, inplace=none}.4, [144])]]

HINT: Re-running with most PyTensor optimizations disabled could provide a back-trace showing when this node was created. This can be done by setting the PyTensor flag 'optimizer=fast_compile'. If that does not work, PyTensor optimizations can be disabled with 'optimizer=None'.
HINT: Use the PyTensor flag exception_verbosity=high for a debug print-out and storage map footprint of this Apply node.

[ricardoV94]

Can you provide a reproducible snippet? The shorter the better

[ricardoV94]

Python code would be better, but can be a NB if that suits you better

[rklees]

Below the code: # Generate dataset measurement_error = st.MeasurementError(name="obs") IRW = st.LevelTrendComponent(order=2, innovations_order=[0, 1]) cycle = st.CycleComponent(name="annual_cycle", cycle_length=12, innovations=True) # cycle length is the period in number of samples; non-integer periods are allowed SA_cycle = st.FrequencySeasonality( name="SA_cycle", season_length=5.347, n=1, innovations=True ) param_dict = { "initial_trend": np.zeros((2,)), "sigma_trend": np.array([0.063]), "annual_cycle": np.array([10., 0.]), "sigma_annual_cycle": np.array([0.316]), "SA_cycle": np.array([20., 0.]), "sigma_SA_cycle": np.array([0.316]), "sigma_obs": np.array([10.]), } mod = IRW + cycle + SA_cycle + measurement_error nobs = 144 x, y = simulate_from_numpy_model(mod, rng, param_dict, steps=nobs) time = np.arange(nobs)/12 data = pd.DataFrame({'time': time, 'meas': y}) # Define model mod = st.LevelTrendComponent(order=2, innovations_order=[0, 1]) mod += st.CycleComponent(name='annual_cycle', cycle_length=12, innovations=True) mod += st.CycleComponent(name='SA_cycle', cycle_length=5.347, innovations=True) mod += st.MeasurementError(name="obs") ssm_model = mod.build(name="IRW + 2 cycles + measurement error") # initial_trend_dims, sigma_trend_dims, annual_cycle_dims, SA_cycle_dims, sigma_obs_dims, P0_dims = ssm_model.param_dims.values() coords = ssm_model.coords # Declare the priors: with pm.Model(coords=coords) as pymc_model: P0_diag = pm.Gamma("P0_diag", alpha=2, beta=5, dims=P0_dims[0]) P0 = pm.Deterministic("P0", pt.diag(P0_diag), dims=P0_dims) initial_trend = pm.Normal("initial_trend", dims=initial_trend_dims) sigma_trend = pm.Gamma("sigma_trend", alpha=2, beta=10, dims=sigma_trend_dims) annual_cycle = pm.Normal("annual_cycle", sigma=5, dims=annual_cycle_dims) sigma_annual_cycle = pm.Gamma("sigma_annual_cycle", alpha=2, beta=5) SA_cycle = pm.Normal("SA_cycle", sigma=5, dims=SA_cycle_dims) sigma_SA_cycle = pm.Gamma("sigma_SA_cycle", alpha=2, beta=5) sigma_obs = pm.Gamma("sigma_obs", alpha=2, beta=5, dims=('observed_state',)) with pymc_model: ssm_model.build_statespace_graph(data['meas'], mode="JAX") samplers = ["pymc", "numpyro", "nutpie", "blackjax"] sampler = samplers[0] with pymc_model: idata = pm.sample(nuts_sampler=sampler, tune=500, draws=1000, chains=4, progressbar=True, target_accept=0.95) From: Ricardo Vieira ***@***.***> Reply to: pymc-devs/pymc-experimental ***@***.***> Date: Thursday, 18 January 2024 at 13:50 To: pymc-devs/pymc-experimental ***@***.***> Cc: Roland Klees ***@***.***>, Author ***@***.***> Subject: Re: [pymc-devs/pymc-experimental] Error messages when using the pymc or nutpie NUTS samplers in combination with pymc-experimental (Issue #298) Can you provide a reproducible snippet? The shorter the better — Reply to this email directly, view it on GitHub<https://urldefense.com/v3/__https:/github.com/pymc-devs/pymc-experimental/issues/298*issuecomment-1898421277__;Iw!!PAKc-5URQlI!4wp-J9U3ZfVtvFqrJz51jkYDftz-xnCxLsA7n_dW7GrH58i2m_fEToCrJiiu0DPL-JgENGzZcRGJskvH_R65LGg_FQ$>, or unsubscribe<https://urldefense.com/v3/__https:/github.com/notifications/unsubscribe-auth/AWKBIZJMYS57HW6L2XHG3PDYPELANAVCNFSM6AAAAABCAG5FZOVHI2DSMVQWIX3LMV43OSLTON2WKQ3PNVWWK3TUHMYTQOJYGQZDCMRXG4__;!!PAKc-5URQlI!4wp-J9U3ZfVtvFqrJz51jkYDftz-xnCxLsA7n_dW7GrH58i2m_fEToCrJiiu0DPL-JgENGzZcRGJskvH_R5bAAdt6g$>. You are receiving this because you authored the thread.Message ID: ***@***.***>

[rklees]

Samplers 'numpyro' and 'blackjax' run smoothly.

[rklees]

Software versions I use:
python version = 3.11.2 (v3.11.2:878ead1ac1, Feb 7 2023, 10:02:41) [Clang 13.0.0 (clang-1300.0.29.30)]
pymc version = 5.10.3
pandas version = 2.1.4
arviz version = 0.17.0
numpy version = 1.26.3
pytensor version = 2.18.6
blackjax version = 1.1.0
nutpie version = 0.9.1
numba version = 0.58.1
xarray version = 2023.12.0

[ricardoV94]

You probably shouldn't set mode="JAX" when building the model if you intend to sample with nutpie. Leave it unchanged or set mode="NUMBA". Nutpie uses the Numba backend, not JAX

[rklees]

This does not solve the problem. Error messages persist.

[jessegrabowski]

This isn't a problem with pymc-experimental per se: nutpie isn't supported for scan-based models. This is a known limitation, and has to do with how random number generators are currently implemented in numba. This issue needs to be closed before it will work.

To use the defaultpymc sampler though, pass mode=None to build_statespace_graph.