Raveled NUTS setup fails for scan model with length-one dimensioned vector RV

[drbenvincent]

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))
# False

Constructing 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

[ricardoV94]

Is it under the limited runtime-broadcasting policy? https://pytensor.readthedocs.io/en/latest/tutorial/broadcasting.html#runtime-broadcasting-limitations

[ricardoV94]

Or a known scan pickiness with shape changes that I think there's a PyTensor issue/PR open for