Description
A PyMC model using pytensor.scan fails during PyMC's in-process NUTS setup when a length-one vector RV with named dims is passed into the scan as a non-sequence.
The model's compile_logp() and compile_dlogp() work, and pm.sample(nuts_sampler="numpyro") works. The failure occurs when PyMC's raveled-input NUTS setup calls join_nonshared_inputs(). Nutpie fails with the same error via an analogous flatten/unflatten replacement path, but this reproduces with nuts_sampler="pymc", so I am reporting it here first.
This originally surfaced downstream in pathmc, but the reproducer below uses only PyMC/PyTensor.
Minimal reproducer
from __future__ import annotations
import numpy as np
import pymc as pm
import pytensor
import pytensor.tensor as pt
n_times, n_units = 8, 3
rng = np.random.default_rng(0)
x = rng.normal(size=(n_times, n_units))
y = 0.7 * np.concatenate([x[0][None, :], x[:-1]], axis=0) + rng.normal(
scale=0.2, size=(n_times, n_units)
)
with pm.Model(coords={"predictors": ["lag(x)"]}) as gen:
x_data = pm.Data("x", x)
lagged_x = pt.concatenate(
[pt.as_tensor_variable(x[0][None, :]), x_data[:-1]], axis=0
)
beta = pm.Normal("beta", 0, 1, dims="predictors")
sigma = pm.HalfNormal("sigma", 1)
def step(x_t, lag_x_t, prev_y, beta):
del x_t, prev_y
return beta[0] * lag_x_t
results = pytensor.scan(
fn=step,
sequences=[x_data, lagged_x],
outputs_info=[pt.as_tensor_variable(np.zeros(n_units, dtype="float64"))],
non_sequences=[beta],
strict=True,
return_updates=False,
)
if not isinstance(results, list):
results = [results]
pm.Deterministic("mu", results[0])
pm.Normal("y", mu=results[0], sigma=sigma, shape=(n_times, n_units))
model = pm.observe(gen, {"y": y})
print("compile_logp", model.compile_logp()(model.initial_point()))
print("compile_dlogp", model.compile_dlogp()(model.initial_point()))
pm.sample(
draws=1,
tune=1,
chains=1,
cores=1,
random_seed=1,
progressbar=False,
compute_convergence_checks=False,
nuts_sampler="pymc",
model=model,
)Observed error
Only 1 samples per chain. Reliable r-hat and ESS diagnostics require longer chains for accurate estimate.
Initializing NUTS using jitter+adapt_diag...
Traceback (most recent call last):
File "<stdin>", line 25, in <module>
File ".../site-packages/pymc/sampling/mcmc.py", line 1022, in sample
initial_points, step = init_nuts(
File ".../site-packages/pymc/sampling/mcmc.py", line 1849, in init_nuts
logp_dlogp_func = model.logp_dlogp_function(ravel_inputs=True, **compile_kwargs)
File ".../site-packages/pymc/model/core.py", line 585, in logp_dlogp_function
return ValueGradFunction(
File ".../site-packages/pymc/model/core.py", line 247, in __init__
outputs, raveled_grad_vars = join_nonshared_inputs(
File ".../site-packages/pymc/pytensorf.py", line 598, in join_nonshared_inputs
new_outputs = [clone_replace(output, replace, rebuild_strict=False) for output in outputs]
File ".../site-packages/pytensor/scan/op.py", line 1267, in make_node
raise ValueError(
ValueError: Argument Reshape{1}.0 given to the scan node is not compatible with its corresponding loop function variable i3
Backend comparison
Using the same model:
| Path |
Result |
model.compile_logp() |
OK |
model.compile_dlogp() |
OK |
pm.sample(nuts_sampler="numpyro") |
OK |
pm.sample(nuts_sampler="pymc") |
Fails with the error above |
nutpie.compile_pymc_model(model) |
Fails with the same scan compatibility error |
pm.sample(nuts_sampler="nutpie") |
Fails with the same scan compatibility error |
Diagnosis
The issue appears to be the replacement created by the raveled-input path.
The original value variable for beta has non-broadcastable unknown-length vector type:
beta Vector(float64, shape=(?,), broadcastable=(False,))
join_nonshared_inputs() builds a replacement from the raveled point via:
joined_inputs[last_idx : last_idx + arr_len].reshape(shape).astype(var.dtype)
For a length-one vector, this replacement becomes:
Reshape{1}.0 Vector(float64, shape=(1,), broadcastable=(True,))
That replacement is not in the same type class as the scan inner non-sequence variable, so pytensor.scan.op.Scan.make_node() rejects it.
A direct PyTensor check shows the mismatch:
import pytensor.tensor as pt
old = pt.vector("old")
joined = pt.dvector("joined")
repl = joined[:1].reshape((1,))
filtered = old.type.filter_variable(repl)
print(old.type, old.type.shape, old.type.broadcastable)
# Vector(float64, shape=(?,)) (None,) (False,)
print(repl.type, repl.type.shape, repl.type.broadcastable)
# Vector(float64, shape=(1,)) (1,) (True,)
print(filtered.type, filtered.type.shape, filtered.type.broadcastable)
# Vector(float64, shape=(1,)) (1,) (True,)
print(repl.type.in_same_class(old.type))
# FalseConstructing a replacement that preserves a non-broadcastable vector axis avoids the direct clone_replace failure, so the fix may be for join_nonshared_inputs() to preserve the original input variable's broadcastability/type when unflattening length-one vector inputs.
Expected behavior
PyMC's in-process NUTS setup should be able to sample this model, or at least the raveled-input replacement should preserve the original value variable's type class so scan non-sequence compatibility is maintained.
Environment
Tested with:
pymc 6.0.1
pytensor 3.0.4
nutpie 0.16.10
numpyro 0.21.0
python 3.12.9
macOS Darwin
Description
A PyMC model using
pytensor.scanfails during PyMC's in-process NUTS setup when a length-one vector RV with nameddimsis passed into the scan as a non-sequence.The model's
compile_logp()andcompile_dlogp()work, andpm.sample(nuts_sampler="numpyro")works. The failure occurs when PyMC's raveled-input NUTS setup callsjoin_nonshared_inputs(). Nutpie fails with the same error via an analogous flatten/unflatten replacement path, but this reproduces withnuts_sampler="pymc", so I am reporting it here first.This originally surfaced downstream in
pathmc, but the reproducer below uses only PyMC/PyTensor.Minimal reproducer
Observed error
Backend comparison
Using the same model:
model.compile_logp()model.compile_dlogp()pm.sample(nuts_sampler="numpyro")pm.sample(nuts_sampler="pymc")nutpie.compile_pymc_model(model)pm.sample(nuts_sampler="nutpie")Diagnosis
The issue appears to be the replacement created by the raveled-input path.
The original value variable for
betahas non-broadcastable unknown-length vector type:join_nonshared_inputs()builds a replacement from the raveled point via:For a length-one vector, this replacement becomes:
That replacement is not in the same type class as the scan inner non-sequence variable, so
pytensor.scan.op.Scan.make_node()rejects it.A direct PyTensor check shows the mismatch:
Constructing a replacement that preserves a non-broadcastable vector axis avoids the direct
clone_replacefailure, so the fix may be forjoin_nonshared_inputs()to preserve the original input variable's broadcastability/type when unflattening length-one vector inputs.Expected behavior
PyMC's in-process NUTS setup should be able to sample this model, or at least the raveled-input replacement should preserve the original value variable's type class so
scannon-sequence compatibility is maintained.Environment
Tested with: