Issues with dimensions of unobserved components in versions 0.1.0 and 0.1.1

[rklees]

I run my standard test example using version 0.1.0 and get an error message from pm.sample, see below. Any advice of how to proceed?

TypeError Traceback (most recent call last)
Cell In[19], line 4
2 sampler = 'numpyro'
3 with pymc_model:
----> 4 idata = pm.sample(nuts_sampler=sampler, tune=500, draws=1000, chains=4, progressbar=True, target_accept=0.95)
6 # idate is an "inference data object", provided by the sampler. Sampling statistics are provided in idata.sample_stats. For more information
7 # about the sampling statistics, see
8 # https://www.pymc.io/projects/docs/en/v3/pymc-examples/examples/diagnostics_and_criticism/sampler-stats.html

File ~/miniforge3/envs/BSSM/lib/python3.11/site-packages/pymc/sampling/mcmc.py:691, in sample(draws, tune, chains, cores, random_seed, progressbar, progressbar_theme, step, var_names, 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)
687 if not isinstance(step, NUTS):
688 raise ValueError(
689 "Model can not be sampled with NUTS alone. Your model is probably not continuous."
690 )
--> 691 return _sample_external_nuts(
692 sampler=nuts_sampler,
693 draws=draws,
694 tune=tune,
695 chains=chains,
696 target_accept=kwargs.pop("nuts", {}).get("target_accept", 0.8),
697 random_seed=random_seed,
698 initvals=initvals,
699 model=model,
700 var_names=var_names,
701 progressbar=progressbar,
702 idata_kwargs=idata_kwargs,
703 nuts_sampler_kwargs=nuts_sampler_kwargs,
704 **kwargs,
705 )
707 if isinstance(step, list):
708 step = CompoundStep(step)

File ~/miniforge3/envs/BSSM/lib/python3.11/site-packages/pymc/sampling/mcmc.py:351, in _sample_external_nuts(sampler, draws, tune, chains, target_accept, random_seed, initvals, model, var_names, progressbar, idata_kwargs, nuts_sampler_kwargs, **kwargs)
348 elif sampler in ("numpyro", "blackjax"):
349 import pymc.sampling.jax as pymc_jax
--> 351 idata = pymc_jax.sample_jax_nuts(
352 draws=draws,
353 tune=tune,
354 chains=chains,
355 target_accept=target_accept,
356 random_seed=random_seed,
357 initvals=initvals,
358 model=model,
359 var_names=var_names,
360 progressbar=progressbar,
361 nuts_sampler=sampler,
362 idata_kwargs=idata_kwargs,
363 **nuts_sampler_kwargs,
364 )
365 return idata
367 else:

File ~/miniforge3/envs/BSSM/lib/python3.11/site-packages/pymc/sampling/jax.py:567, in sample_jax_nuts(draws, tune, chains, target_accept, random_seed, initvals, jitter, model, var_names, nuts_kwargs, progressbar, keep_untransformed, chain_method, postprocessing_backend, postprocessing_vectorize, postprocessing_chunks, idata_kwargs, compute_convergence_checks, nuts_sampler)
564 raise ValueError(f"{nuts_sampler=} not recognized")
566 tic1 = datetime.now()
--> 567 raw_mcmc_samples, sample_stats, library = sampler_fn(
568 model=model,
569 target_accept=target_accept,
570 tune=tune,
571 draws=draws,
572 chains=chains,
573 chain_method=chain_method,
574 progressbar=progressbar,
575 random_seed=random_seed,
576 initial_points=initial_points,
577 nuts_kwargs=nuts_kwargs,
578 )
579 tic2 = datetime.now()
581 jax_fn = get_jaxified_graph(inputs=model.value_vars, outputs=vars_to_sample)

File ~/miniforge3/envs/BSSM/lib/python3.11/site-packages/pymc/sampling/jax.py:458, in _sample_numpyro_nuts(model, target_accept, tune, draws, chains, chain_method, progressbar, random_seed, initial_points, nuts_kwargs)
454 import numpyro
456 from numpyro.infer import MCMC, NUTS
--> 458 logp_fn = get_jaxified_logp(model, negative_logp=False)
460 nuts_kwargs.setdefault("adapt_step_size", True)
461 nuts_kwargs.setdefault("adapt_mass_matrix", True)

File ~/miniforge3/envs/BSSM/lib/python3.11/site-packages/pymc/sampling/jax.py:153, in get_jaxified_logp(model, negative_logp)
151 if not negative_logp:
152 model_logp = -model_logp
--> 153 logp_fn = get_jaxified_graph(inputs=model.value_vars, outputs=[model_logp])
155 def logp_fn_wrap(x):
156 return logp_fn(*x)[0]

File ~/miniforge3/envs/BSSM/lib/python3.11/site-packages/pymc/sampling/jax.py:128, in get_jaxified_graph(inputs, outputs)
122 def get_jaxified_graph(
123 inputs: list[TensorVariable] | None = None,
124 outputs: list[TensorVariable] | None = None,
125 ) -> list[TensorVariable]:
126 """Compile an PyTensor graph into an optimized JAX function"""
--> 128 graph = _replace_shared_variables(outputs) if outputs is not None else None
130 fgraph = FunctionGraph(inputs=inputs, outputs=graph, clone=True)
131 # We need to add a Supervisor to the fgraph to be able to run the
132 # JAX sequential optimizer without warnings. We made sure there
133 # are no mutable input variables, so we only need to check for
134 # "destroyers". This should be automatically handled by PyTensor
135 # once aesara-devs/aesara#637 is fixed.

File ~/miniforge3/envs/BSSM/lib/python3.11/site-packages/pymc/sampling/jax.py:118, in _replace_shared_variables(graph)
111 raise ValueError(
112 "Graph contains shared variables with default_update which cannot "
113 "be safely replaced."
114 )
116 replacements = {var: pt.constant(var.get_value(borrow=True)) for var in shared_variables}
--> 118 new_graph = clone_replace(graph, replace=replacements)
119 return new_graph

File ~/miniforge3/envs/BSSM/lib/python3.11/site-packages/pytensor/graph/replace.py:85, in clone_replace(output, replace, **rebuild_kwds)
82 _, _outs, _ = rebuild_collect_shared(output, [], tmp_replace, [], **rebuild_kwds)
84 # TODO Explain why we call it twice ?!
---> 85 _, outs, _ = rebuild_collect_shared(_outs, [], new_replace, [], **rebuild_kwds)
87 return outs

File ~/miniforge3/envs/BSSM/lib/python3.11/site-packages/pytensor/compile/function/pfunc.py:313, in rebuild_collect_shared(outputs, inputs, replace, updates, rebuild_strict, copy_inputs_over, no_default_updates, clone_inner_graphs)
311 for v in outputs:
312 if isinstance(v, Variable):
--> 313 cloned_v = clone_v_get_shared_updates(v, copy_inputs_over)
314 cloned_outputs.append(cloned_v)
315 elif isinstance(v, Out):

File ~/miniforge3/envs/BSSM/lib/python3.11/site-packages/pytensor/compile/function/pfunc.py:189, in rebuild_collect_shared..clone_v_get_shared_updates(v, copy_inputs_over)
187 if owner not in clone_d:
188 for i in owner.inputs:
--> 189 clone_v_get_shared_updates(i, copy_inputs_over)
190 clone_node_and_cache(
191 owner,
192 clone_d,
193 strict=rebuild_strict,
194 clone_inner_graphs=clone_inner_graphs,
195 )
196 return clone_d.setdefault(v, v)

File ~/miniforge3/envs/BSSM/lib/python3.11/site-packages/pytensor/compile/function/pfunc.py:189, in rebuild_collect_shared..clone_v_get_shared_updates(v, copy_inputs_over)
187 if owner not in clone_d:
188 for i in owner.inputs:
--> 189 clone_v_get_shared_updates(i, copy_inputs_over)
190 clone_node_and_cache(
191 owner,
192 clone_d,
193 strict=rebuild_strict,
194 clone_inner_graphs=clone_inner_graphs,
195 )
196 return clone_d.setdefault(v, v)

[... skipping similar frames: rebuild_collect_shared.<locals>.clone_v_get_shared_updates at line 189 (2 times)]

File ~/miniforge3/envs/BSSM/lib/python3.11/site-packages/pytensor/compile/function/pfunc.py:189, in rebuild_collect_shared..clone_v_get_shared_updates(v, copy_inputs_over)
187 if owner not in clone_d:
188 for i in owner.inputs:
--> 189 clone_v_get_shared_updates(i, copy_inputs_over)
190 clone_node_and_cache(
191 owner,
192 clone_d,
193 strict=rebuild_strict,
194 clone_inner_graphs=clone_inner_graphs,
195 )
196 return clone_d.setdefault(v, v)

File ~/miniforge3/envs/BSSM/lib/python3.11/site-packages/pytensor/compile/function/pfunc.py:190, in rebuild_collect_shared..clone_v_get_shared_updates(v, copy_inputs_over)
188 for i in owner.inputs:
189 clone_v_get_shared_updates(i, copy_inputs_over)
--> 190 clone_node_and_cache(
191 owner,
192 clone_d,
193 strict=rebuild_strict,
194 clone_inner_graphs=clone_inner_graphs,
195 )
196 return clone_d.setdefault(v, v)
197 elif isinstance(v, SharedVariable):

File ~/miniforge3/envs/BSSM/lib/python3.11/site-packages/pytensor/graph/basic.py:1201, in clone_node_and_cache(node, clone_d, clone_inner_graphs, **kwargs)
1197 new_op: "Op" | None = cast(Optional["Op"], clone_d.get(node.op))
1199 cloned_inputs: list[Variable] = [cast(Variable, clone_d[i]) for i in node.inputs]
-> 1201 new_node = node.clone_with_new_inputs(
1202 cloned_inputs,
1203 # Only clone inner-graph Ops when there isn't a cached clone (and
1204 # when clone_inner_graphs is enabled)
1205 clone_inner_graph=clone_inner_graphs if new_op is None else False,
1206 **kwargs,
1207 )
1209 if new_op:
1210 # If we didn't clone the inner-graph Op above, because
1211 # there was a cached version, set the cloned Apply to use
1212 # the cached clone Op
1213 new_node.op = new_op

File ~/miniforge3/envs/BSSM/lib/python3.11/site-packages/pytensor/graph/basic.py:285, in Apply.clone_with_new_inputs(self, inputs, strict, clone_inner_graph)
282 if isinstance(new_op, HasInnerGraph) and clone_inner_graph: # type: ignore
283 new_op = new_op.clone() # type: ignore
--> 285 new_node = new_op.make_node(*new_inputs)
286 new_node.tag = copy(self.tag).update(new_node.tag)
287 else:

File ~/miniforge3/envs/BSSM/lib/python3.11/site-packages/pytensor/scan/op.py:964, in Scan.make_node(self, *inputs)
960 argoffset = 0
961 for inner_seq, outer_seq in zip(
962 self.inner_seqs(self.inner_inputs), self.outer_seqs(inputs)
963 ):
--> 964 check_broadcast(outer_seq, inner_seq)
965 new_inputs.append(copy_var_format(outer_seq, as_var=inner_seq))
967 argoffset += len(self.outer_seqs(inputs))

File ~/miniforge3/envs/BSSM/lib/python3.11/site-packages/pytensor/scan/op.py:179, in check_broadcast(v1, v2)
177 a1 = n + size - v1.type.ndim + 1
178 a2 = n + size - v2.type.ndim + 1
--> 179 raise TypeError(msg % (v1.type, v2.type, a1, b1, b2, a2))

TypeError: The broadcast pattern of the output of scan (Matrix(float64, shape=(144, 1))) is inconsistent with the one provided in output_info (Vector(float64, shape=(?,))). The output on axis 0 is True, but it is False on axis 1 in output_info. This can happen if one of the dimension is fixed to 1 in the input, while it is still variable in the output, or vice-verca. You have to make them consistent, e.g. using pytensor.tensor.{unbroadcast, specify_broadcastable}.

[ricardoV94]

You're probably specifying a coords of length 1 which in the new version are mutable by default. You can specify_broadcastable to inform PyMC that this dimension will always be 1 or not specify it and instead use expand_dims, so PyMC also knows this will always have length 1.

[rklees]

I don't specify any coords. I set up a model as
IRW = st.LevelTrendComponent(order=2, innovations_order=[0, 1])
annual_cycle = st.CycleComponent(name='annual_cycle', cycle_length=12, innovations=True)
SA_cycle = st.CycleComponent(name='SA_cycle', cycle_length=5.347, innovations=True)
obs_noise = st.MeasurementError(name="obs")
pymc_mod = (IRW + annual_cycle + SA_cycle + obs_noise).build(name="IRW + 2 cycles + measurement error")

initial_trend_dims, sigma_trend_dims, annual_cycle_dims, SA_cycle_dims, P0_dims = pymc_mod.param_dims.values()

define the priors, build the graph and sample

with pm.Model(coords=coords) as pymc_model:
# priors for initial state vector comprising initial_trend (i.e., level and slope), initial annual cycle and initial SA cycle
initial_trend = pm.Normal("initial_trend", sigma=100, dims=initial_trend_dims)
initial_annual_cycle = pm.Normal("annual_cycle", shape=(2,), sigma=100, dims=annual_cycle_dims)
initial_SA_cycle = pm.Normal("SA_cycle", shape=(2,), sigma=100, dims=SA_cycle_dims)
# priors for square root of initial state covariance
P0_diag = pm.Gamma("P0_diag", alpha=2, beta=5)
P0 = pm.Deterministic("P0", pt.eye(pymc_mod.k_states) * P0_diag, dims=P0_dims)
# Priors for model parameters
sigma_trend = pm.Gamma("sigma_trend", alpha=2, beta=10, dims=sigma_trend_dims)
sigma_annual_cycle = pm.Gamma("sigma_annual_cycle", alpha=2, beta=5)
sigma_SA_cycle = pm.Gamma("sigma_SA_cycle", alpha=2, beta=5)
sigma_obs = pm.HalfNormal("sigma_obs", sigma=10.)
SA_cycle_length = pm.HalfNormal("SA_cycle_length", sigma=5)

pymc_mod.build_statespace_graph(pd.DataFrame(data['meas']), mode="JAX")
idata = pm.sample(nuts_sampler=sampler, tune=500, draws=1000, chains=4, progressbar=True, target_accept=0.95)

Any specific help would be very much appreciated. Thanks.

[jessegrabowski]

Looks like another thing broken by pymc-devs/pymc#7047. I will look at it ASAP.

Thanks for reporting all these bugs by the way, it's really important we get them fixed. I deeply appreciate your effort.

[rklees]

Any update regarding this issue?

[rklees]

This problem still persists in version 0.1.1. Is there any hope that someone will fix it? Otherwise, I can't use pymc-experimental for unobserved components. Would be a pitty.

[rklees]

Here are the details:
NotImplementedError Traceback (most recent call last)
File ~/miniforge3/envs/BSSM/lib/python3.11/site-packages/pytensor/link/basic.py:102, in Container.set(self, value)
100 try:
101 # Use in-place filtering when/if possible
--> 102 self.storage[0] = self.type.filter_inplace(
103 value, self.storage[0], **kwargs
104 )
105 except NotImplementedError:

File ~/miniforge3/envs/BSSM/lib/python3.11/site-packages/pytensor/graph/type.py:128, in Type.filter_inplace(self, value, storage, strict, allow_downcast)
109 """Return data or an appropriately wrapped/converted data by converting it in-place.
110
111 This method allows one to reuse old allocated memory. If this method
(...)
126 NotImplementedError
127 """
--> 128 raise NotImplementedError()

NotImplementedError:

During handling of the above exception, another exception occurred:

TypeError Traceback (most recent call last)
Cell In[5], line 26
23 nobs = 144 # we simulate a time series of length Nobs
25 # simulate data (by simulating each SSM component and sum up)
---> 26 x, y = simulate_from_numpy_model(mod, rng, param_dict, steps=nobs)
27 print(type(y))
29 # we standardize the data
30 # y = (y - np.mean(y))/np.std(y)

Cell In[3], line 17, in simulate_from_numpy_model(mod, rng, param_dict, steps)
13 def simulate_from_numpy_model(mod, rng, param_dict, steps=100):
14 """
15 Helper function to visualize the components outside of a PyMC model context
16 """
---> 17 x0, P0, c, d, T, Z, R, H, Q = unpack_symbolic_matrices_with_params(mod, param_dict)
18 Z_time_varies = Z.ndim == 3
20 k_states = mod.k_states

Cell In[3], line 9, in unpack_symbolic_matrices_with_params(mod, param_dict)
5 def unpack_symbolic_matrices_with_params(mod, param_dict):
6 f_matrices = pytensor.function(
7 list(mod._name_to_variable.values()), unpack_statespace(mod.ssm), on_unused_input="ignore"
8 )
----> 9 x0, P0, c, d, T, Z, R, H, Q = f_matrices(**param_dict)
10 return x0, P0, c, d, T, Z, R, H, Q

File ~/miniforge3/envs/BSSM/lib/python3.11/site-packages/pytensor/compile/function/types.py:897, in Function.call(self, *args, **kwargs)
895 if kwargs: # for speed, skip the items for empty kwargs
896 for k, arg in kwargs.items():
--> 897 self[k] = arg
899 if (
900 not self.trust_input
901 and
(...)
904 ):
905 # Collect aliased inputs among the storage space
906 args_share_memory = []

File ~/miniforge3/envs/BSSM/lib/python3.11/site-packages/pytensor/compile/function/types.py:549, in Function.setitem(self, item, value)
548 def setitem(self, item, value):
--> 549 self.value[item] = value

File ~/miniforge3/envs/BSSM/lib/python3.11/site-packages/pytensor/compile/function/types.py:505, in Function.init..ValueAttribute.setitem(self, item, value)
499 raise TypeError(
500 f"Ambiguous name: {item} - please check the "
501 "names of the inputs of your function "
502 "for duplicates."
503 )
504 if isinstance(s, Container):
--> 505 s.value = value
506 s.provided += 1
507 else:

File ~/miniforge3/envs/BSSM/lib/python3.11/site-packages/pytensor/link/basic.py:106, in Container.set(self, value)
102 self.storage[0] = self.type.filter_inplace(
103 value, self.storage[0], **kwargs
104 )
105 except NotImplementedError:
--> 106 self.storage[0] = self.type.filter(value, **kwargs)
108 except Exception as e:
109 e.args = (*e.args, f'Container name "{self.name}"')

File ~/miniforge3/envs/BSSM/lib/python3.11/site-packages/pytensor/tensor/type.py:242, in TensorType.filter(self, data, strict, allow_downcast)
239 raise TypeError(err_msg)
241 if self.ndim != data.ndim:
--> 242 raise TypeError(
243 f"Wrong number of dimensions: expected {self.ndim},"
244 f" got {data.ndim} with shape {data.shape}."
245 )
246 if not data.flags.aligned:
247 raise TypeError(
248 "The numpy.ndarray object is not aligned."
249 " PyTensor C code does not support that.",
250 )

TypeError: ('Wrong number of dimensions: expected 0, got 1 with shape (1,).', 'Container name "sigma_annual_cycle"')

[jessegrabowski]

#346 should address this

[rklees]

How to get a version with this issue fixed?

[rklees]

I will use the latest version of the example notebook and do some tests.

[rklees]

Two issues: 1. I get an error message when running the notebook “Structural Timeseries Modeling.ipynb" 2. When adding a cycle, I still get the same error message as mentioned in my first post. I run the following packages: pymc version = 5.13.1 pytensor version = 2.20.0 pandas version = 2.2.1 arviz version = 0.18.0 numpy version = 1.26.3 pytensor version = 2.20.0 blackjax version = 0.0.0 nutpie version = 0.10.0 numba version = 0.58.1 xarray version = 2024.3.0 pymc-experimental = 0.1.1 The error message when running the notebook is

…

--------------------------------------------------------------------------- TypeError Traceback (most recent call last) Cell In[34], line 13 10 sigma_obs = pm.HalfNormal("sigma_obs", sigma=0.05) 12 ss_mod.build_statespace_graph(nile, mode="JAX")

---> 13 idata = pm.sample(nuts_sampler="numpyro", chains=8, draws=500) File ~/miniforge3/envs/BSSM/lib/python3.11/site-packages/pymc/sampling/mcmc.py:691<http://localhost:8889/lab/tree/Software/Python%20software/PyMC%20experimental%20Notebooks/~/miniforge3/envs/BSSM/lib/python3.11/site-packages/pymc/sampling/mcmc.py#line=690>, in sample(draws, tune, chains, cores, random_seed, progressbar, progressbar_theme, step, var_names, 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) 687 if not isinstance(step, NUTS): 688 raise ValueError( 689 "Model can not be sampled with NUTS alone. Your model is probably not continuous." 690 )

--> 691 return _sample_external_nuts( 692 sampler=nuts_sampler, 693 draws=draws, 694 tune=tune, 695 chains=chains, 696 target_accept=kwargs.pop("nuts", {}).get("target_accept", 0.8), 697 random_seed=random_seed, 698 initvals=initvals, 699 model=model, 700 var_names=var_names, 701 progressbar=progressbar, 702 idata_kwargs=idata_kwargs, 703 nuts_sampler_kwargs=nuts_sampler_kwargs, 704 **kwargs, 705 ) 707 if isinstance(step, list): 708 step = CompoundStep(step) File ~/miniforge3/envs/BSSM/lib/python3.11/site-packages/pymc/sampling/mcmc.py:351<http://localhost:8889/lab/tree/Software/Python%20software/PyMC%20experimental%20Notebooks/~/miniforge3/envs/BSSM/lib/python3.11/site-packages/pymc/sampling/mcmc.py#line=350>, in _sample_external_nuts(sampler, draws, tune, chains, target_accept, random_seed, initvals, model, var_names, progressbar, idata_kwargs, nuts_sampler_kwargs, **kwargs) 348 elif sampler in ("numpyro", "blackjax"): 349 import pymc.sampling.jax as pymc_jax

--> 351 idata = pymc_jax.sample_jax_nuts( 352 draws=draws, 353 tune=tune, 354 chains=chains, 355 target_accept=target_accept, 356 random_seed=random_seed, 357 initvals=initvals, 358 model=model, 359 var_names=var_names, 360 progressbar=progressbar, 361 nuts_sampler=sampler, 362 idata_kwargs=idata_kwargs, 363 **nuts_sampler_kwargs, 364 ) 365 return idata 367 else: File ~/miniforge3/envs/BSSM/lib/python3.11/site-packages/pymc/sampling/jax.py:567<http://localhost:8889/lab/tree/Software/Python%20software/PyMC%20experimental%20Notebooks/~/miniforge3/envs/BSSM/lib/python3.11/site-packages/pymc/sampling/jax.py#line=566>, in sample_jax_nuts(draws, tune, chains, target_accept, random_seed, initvals, jitter, model, var_names, nuts_kwargs, progressbar, keep_untransformed, chain_method, postprocessing_backend, postprocessing_vectorize, postprocessing_chunks, idata_kwargs, compute_convergence_checks, nuts_sampler) 564 raise ValueError(f"{nuts_sampler=} not recognized") 566 tic1 = datetime.now()

--> 567 raw_mcmc_samples, sample_stats, library = sampler_fn( 568 model=model, 569 target_accept=target_accept, 570 tune=tune, 571 draws=draws, 572 chains=chains, 573 chain_method=chain_method, 574 progressbar=progressbar, 575 random_seed=random_seed, 576 initial_points=initial_points, 577 nuts_kwargs=nuts_kwargs, 578 ) 579 tic2 = datetime.now() 581 jax_fn = get_jaxified_graph(inputs=model.value_vars, outputs=vars_to_sample) File ~/miniforge3/envs/BSSM/lib/python3.11/site-packages/pymc/sampling/jax.py:458<http://localhost:8889/lab/tree/Software/Python%20software/PyMC%20experimental%20Notebooks/~/miniforge3/envs/BSSM/lib/python3.11/site-packages/pymc/sampling/jax.py#line=457>, in _sample_numpyro_nuts(model, target_accept, tune, draws, chains, chain_method, progressbar, random_seed, initial_points, nuts_kwargs) 454 import numpyro 456 from numpyro.infer import MCMC, NUTS

--> 458 logp_fn = get_jaxified_logp(model, negative_logp=False) 460 nuts_kwargs.setdefault("adapt_step_size", True) 461 nuts_kwargs.setdefault("adapt_mass_matrix", True) File ~/miniforge3/envs/BSSM/lib/python3.11/site-packages/pymc/sampling/jax.py:153<http://localhost:8889/lab/tree/Software/Python%20software/PyMC%20experimental%20Notebooks/~/miniforge3/envs/BSSM/lib/python3.11/site-packages/pymc/sampling/jax.py#line=152>, in get_jaxified_logp(model, negative_logp) 151 if not negative_logp: 152 model_logp = -model_logp

--> 153 logp_fn = get_jaxified_graph(inputs=model.value_vars, outputs=[model_logp]) 155 def logp_fn_wrap(x): 156 return logp_fn(*x)[0] File ~/miniforge3/envs/BSSM/lib/python3.11/site-packages/pymc/sampling/jax.py:128<http://localhost:8889/lab/tree/Software/Python%20software/PyMC%20experimental%20Notebooks/~/miniforge3/envs/BSSM/lib/python3.11/site-packages/pymc/sampling/jax.py#line=127>, in get_jaxified_graph(inputs, outputs) 122 def get_jaxified_graph( 123 inputs: list[TensorVariable] | None = None, 124 outputs: list[TensorVariable] | None = None, 125 ) -> list[TensorVariable]: 126 """Compile an PyTensor graph into an optimized JAX function"""

--> 128 graph = _replace_shared_variables(outputs) if outputs is not None else None 130 fgraph = FunctionGraph(inputs=inputs, outputs=graph, clone=True) 131 # We need to add a Supervisor to the fgraph to be able to run the 132 # JAX sequential optimizer without warnings. We made sure there 133 # are no mutable input variables, so we only need to check for 134 # "destroyers". This should be automatically handled by PyTensor 135 # once aesara-devs/aesara#637 is fixed. File ~/miniforge3/envs/BSSM/lib/python3.11/site-packages/pymc/sampling/jax.py:118, in _replace_shared_variables(graph) 111 raise ValueError( 112 "Graph contains shared variables with default_update which cannot " 113 "be safely replaced." 114 ) 116 replacements = {var: pt.constant(var.get_value(borrow=True)) for var in shared_variables}

--> 118 new_graph = clone_replace(graph, replace=replacements) 119 return new_graph File ~/miniforge3/envs/BSSM/lib/python3.11/site-packages/pytensor/graph/replace.py:85<http://localhost:8889/lab/tree/Software/Python%20software/PyMC%20experimental%20Notebooks/~/miniforge3/envs/BSSM/lib/python3.11/site-packages/pytensor/graph/replace.py#line=84>, in clone_replace(output, replace, **rebuild_kwds) 82 _, _outs, _ = rebuild_collect_shared(output, [], tmp_replace, [], **rebuild_kwds) 84 # TODO Explain why we call it twice ?!

---> 85 _, outs, _ = rebuild_collect_shared(_outs, [], new_replace, [], **rebuild_kwds) 87 return outs File ~/miniforge3/envs/BSSM/lib/python3.11/site-packages/pytensor/compile/function/pfunc.py:313<http://localhost:8889/lab/tree/Software/Python%20software/PyMC%20experimental%20Notebooks/~/miniforge3/envs/BSSM/lib/python3.11/site-packages/pytensor/compile/function/pfunc.py#line=312>, in rebuild_collect_shared(outputs, inputs, replace, updates, rebuild_strict, copy_inputs_over, no_default_updates, clone_inner_graphs) 311 for v in outputs: 312 if isinstance(v, Variable):

--> 313 cloned_v = clone_v_get_shared_updates(v, copy_inputs_over) 314 cloned_outputs.append(cloned_v) 315 elif isinstance(v, Out): File ~/miniforge3/envs/BSSM/lib/python3.11/site-packages/pytensor/compile/function/pfunc.py:189<http://localhost:8889/lab/tree/Software/Python%20software/PyMC%20experimental%20Notebooks/~/miniforge3/envs/BSSM/lib/python3.11/site-packages/pytensor/compile/function/pfunc.py#line=188>, in rebuild_collect_shared.<locals>.clone_v_get_shared_updates(v, copy_inputs_over) 187 if owner not in clone_d: 188 for i in owner.inputs:

--> 189 clone_v_get_shared_updates(i, copy_inputs_over) 190 clone_node_and_cache( 191 owner, 192 clone_d, 193 strict=rebuild_strict, 194 clone_inner_graphs=clone_inner_graphs, 195 ) 196 return clone_d.setdefault(v, v) File ~/miniforge3/envs/BSSM/lib/python3.11/site-packages/pytensor/compile/function/pfunc.py:189<http://localhost:8889/lab/tree/Software/Python%20software/PyMC%20experimental%20Notebooks/~/miniforge3/envs/BSSM/lib/python3.11/site-packages/pytensor/compile/function/pfunc.py#line=188>, in rebuild_collect_shared.<locals>.clone_v_get_shared_updates(v, copy_inputs_over) 187 if owner not in clone_d: 188 for i in owner.inputs:

--> 189 clone_v_get_shared_updates(i, copy_inputs_over) 190 clone_node_and_cache( 191 owner, 192 clone_d, 193 strict=rebuild_strict, 194 clone_inner_graphs=clone_inner_graphs, 195 ) 196 return clone_d.setdefault(v, v) [... skipping similar frames: rebuild_collect_shared.<locals>.clone_v_get_shared_updates at line 189 (2 times)] File ~/miniforge3/envs/BSSM/lib/python3.11/site-packages/pytensor/compile/function/pfunc.py:189<http://localhost:8889/lab/tree/Software/Python%20software/PyMC%20experimental%20Notebooks/~/miniforge3/envs/BSSM/lib/python3.11/site-packages/pytensor/compile/function/pfunc.py#line=188>, in rebuild_collect_shared.<locals>.clone_v_get_shared_updates(v, copy_inputs_over) 187 if owner not in clone_d: 188 for i in owner.inputs:

--> 189 clone_v_get_shared_updates(i, copy_inputs_over) 190 clone_node_and_cache( 191 owner, 192 clone_d, 193 strict=rebuild_strict, 194 clone_inner_graphs=clone_inner_graphs, 195 ) 196 return clone_d.setdefault(v, v) File ~/miniforge3/envs/BSSM/lib/python3.11/site-packages/pytensor/compile/function/pfunc.py:190<http://localhost:8889/lab/tree/Software/Python%20software/PyMC%20experimental%20Notebooks/~/miniforge3/envs/BSSM/lib/python3.11/site-packages/pytensor/compile/function/pfunc.py#line=189>, in rebuild_collect_shared.<locals>.clone_v_get_shared_updates(v, copy_inputs_over) 188 for i in owner.inputs: 189 clone_v_get_shared_updates(i, copy_inputs_over)

--> 190 clone_node_and_cache( 191 owner, 192 clone_d, 193 strict=rebuild_strict, 194 clone_inner_graphs=clone_inner_graphs, 195 ) 196 return clone_d.setdefault(v, v) 197 elif isinstance(v, SharedVariable): File ~/miniforge3/envs/BSSM/lib/python3.11/site-packages/pytensor/graph/basic.py:1201<http://localhost:8889/lab/tree/Software/Python%20software/PyMC%20experimental%20Notebooks/~/miniforge3/envs/BSSM/lib/python3.11/site-packages/pytensor/graph/basic.py#line=1200>, in clone_node_and_cache(node, clone_d, clone_inner_graphs, **kwargs) 1197 new_op: "Op" | None = cast(Optional["Op"], clone_d.get(node.op)) 1199 cloned_inputs: list[Variable] = [cast(Variable, clone_d[i]) for i in node.inputs] -> 1201 new_node = node.clone_with_new_inputs( 1202 cloned_inputs, 1203 # Only clone inner-graph `Op`s when there isn't a cached clone (and 1204 # when `clone_inner_graphs` is enabled) 1205 clone_inner_graph=clone_inner_graphs if new_op is None else False, 1206 **kwargs, 1207 ) 1209 if new_op: 1210 # If we didn't clone the inner-graph `Op` above, because 1211 # there was a cached version, set the cloned `Apply` to use 1212 # the cached clone `Op` 1213 new_node.op = new_op File ~/miniforge3/envs/BSSM/lib/python3.11/site-packages/pytensor/graph/basic.py:285<http://localhost:8889/lab/tree/Software/Python%20software/PyMC%20experimental%20Notebooks/~/miniforge3/envs/BSSM/lib/python3.11/site-packages/pytensor/graph/basic.py#line=284>, in Apply.clone_with_new_inputs(self, inputs, strict, clone_inner_graph) 282 if isinstance(new_op, HasInnerGraph) and clone_inner_graph: # type: ignore 283 new_op = new_op.clone() # type: ignore

--> 285 new_node = new_op.make_node(*new_inputs) 286 new_node.tag = copy(self.tag).__update__(new_node.tag) 287 else: File ~/miniforge3/envs/BSSM/lib/python3.11/site-packages/pytensor/scan/op.py:964<http://localhost:8889/lab/tree/Software/Python%20software/PyMC%20experimental%20Notebooks/~/miniforge3/envs/BSSM/lib/python3.11/site-packages/pytensor/scan/op.py#line=963>, in Scan.make_node(self, *inputs) 960 argoffset = 0 961 for inner_seq, outer_seq in zip( 962 self.inner_seqs(self.inner_inputs), self.outer_seqs(inputs) 963 ):

--> 964 check_broadcast(outer_seq, inner_seq) 965 new_inputs.append(copy_var_format(outer_seq, as_var=inner_seq)) 967 argoffset += len(self.outer_seqs(inputs)) File ~/miniforge3/envs/BSSM/lib/python3.11/site-packages/pytensor/scan/op.py:179<http://localhost:8889/lab/tree/Software/Python%20software/PyMC%20experimental%20Notebooks/~/miniforge3/envs/BSSM/lib/python3.11/site-packages/pytensor/scan/op.py#line=178>, in check_broadcast(v1, v2) 177 a1 = n + size - v1.type.ndim + 1 178 a2 = n + size - v2.type.ndim + 1

--> 179 raise TypeError(msg % (v1.type, v2.type, a1, b1, b2, a2)) TypeError: The broadcast pattern of the output of scan (Matrix(float64, shape=(100, 1))) is inconsistent with the one provided in `output_info` (Vector(float64, shape=(?,))). The output on axis 0 is `True`, but it is `False` on axis 1 in `output_info`. This can happen if one of the dimension is fixed to 1 in the input, while it is still variable in the output, or vice-verca. You have to make them consistent, e.g. using pytensor.tensor.{unbroadcast, specify_broadcastable}.

[rklees]

The error message in your notebook "Structural Time Series Modeling.ipynb" may be related to numpyro. Removing mode=”Jax” in build_statespace_graph and using the pymc sampler works, though sampling is extremely slow. I use numpyro version 0.15.0 and JAX 0.4.27.

[jessegrabowski]

This should be fixed in main now that #346 is merged

[rklees]

Do I just do pip install git+https://github.com/pymc-devs/pymc-experimental.git ??? From: Jesse Grabowski ***@***.***> Reply to: pymc-devs/pymc-experimental ***@***.***> Date: Saturday, 29 June 2024 at 09:56 To: pymc-devs/pymc-experimental ***@***.***> Cc: Roland Klees ***@***.***>, Author ***@***.***> Subject: Re: [pymc-devs/pymc-experimental] Issues with dimensions of unobserved components in versions 0.1.0 and 0.1.1 (Issue #338) This should be fixed in main now that #346<#346> is merged — Reply to this email directly, view it on GitHub<#338 (comment)>, or unsubscribe<https://github.com/notifications/unsubscribe-auth/AWKBIZN44LUDACC56YFK55TZJZSCDAVCNFSM6AAAAABG2JQB5WVHI2DSMVQWIX3LMV43OSLTON2WKQ3PNVWWK3TUHMZDCOJYGAZTMMRUHA>. You are receiving this because you authored the thread.Message ID: ***@***.***>

[zaxtax]

Yes but also try to have the version of pymc on github

…

On Sat, 29 Jun 2024, 15:47 rklees, ***@***.***> wrote: Do I just do pip install git+https://github.com/pymc-devs/pymc-experimental.git ??? From: Jesse Grabowski ***@***.***> Reply to: pymc-devs/pymc-experimental ***@***.***> Date: Saturday, 29 June 2024 at 09:56 To: pymc-devs/pymc-experimental ***@***.***> Cc: Roland Klees ***@***.***>, Author ***@***.***> Subject: Re: [pymc-devs/pymc-experimental] Issues with dimensions of unobserved components in versions 0.1.0 and 0.1.1 (Issue #338) This should be fixed in main now that #346< #346> is merged — Reply to this email directly, view it on GitHub< #338 (comment)>, or unsubscribe< https://github.com/notifications/unsubscribe-auth/AWKBIZN44LUDACC56YFK55TZJZSCDAVCNFSM6AAAAABG2JQB5WVHI2DSMVQWIX3LMV43OSLTON2WKQ3PNVWWK3TUHMZDCOJYGAZTMMRUHA>. You are receiving this because you authored the thread.Message ID: ***@***.***> — Reply to this email directly, view it on GitHub <#338 (comment)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AAACCULUM5255ZAGCFCJDW3ZJ23GTAVCNFSM6AAAAABG2JQB5WVHI2DSMVQWIX3LMV43OSLTON2WKQ3PNVWWK3TUHMZDCOJYGIYDAMZRGE> . You are receiving this because you are subscribed to this thread.Message ID: ***@***.***>