BUG: Adding a Deterministic results in the error: All variables needed to compute inner-graph must be provided as inputs under strict=True

[tomicapretto]

Describe the issue:

The same model results in an error depending on the usage of a deterministic. See the examples below.

Reproduceable code example:

import numpy as np
import pymc as pm
import pytensor.tensor as pt

rng = np.random.default_rng(1234)
n1, n2 = 30, 70
y = np.concatenate([np.zeros(n1), rng.poisson(3, size=n2)]).astype(int)

coords={"__obs__": np.arange(n1 + n2)}


# Works
with pm.Model(coords=coords) as model:
    a = pm.Normal("Intercept", mu=0, sigma=2.5)
    psi = pm.Beta("psi", alpha=2, beta=2)
    pm.HurdlePoisson("y", mu=pt.exp(a), psi=psi, observed=y, dims="__obs__")

# Raises error
with pm.Model(coords=coords) as model:
    a = pm.Normal("Intercept", mu=0, sigma=2.5)
    mu = pm.Deterministic("mu", pt.exp(a))
    psi = pm.Beta("psi", alpha=2, beta=2)
    pm.HurdlePoisson("y", mu=mu, psi=psi, observed=y, dims="__obs__")

Error message:

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
Cell In[3], line 5
      3 mu = pm.Deterministic("mu", pt.exp(a))
      4 psi = pm.Beta("psi", alpha=2, beta=2)
----> 5 pm.HurdlePoisson("y", mu=mu, psi=psi, observed=y, dims="__obs__")

File ~/anaconda3/envs/bambi-dev/lib/python3.11/site-packages/pymc/distributions/mixture.py:866, in HurdlePoisson.__new__(cls, name, psi, mu, **kwargs)
    865 def __new__(cls, name, psi, mu, **kwargs):
--> 866     return _hurdle_mixture(
    867         name=name, nonzero_p=psi, nonzero_dist=Poisson.dist(mu=mu), dtype="int", **kwargs
    868     )

File ~/anaconda3/envs/bambi-dev/lib/python3.11/site-packages/pymc/distributions/mixture.py:822, in _hurdle_mixture(name, nonzero_p, nonzero_dist, dtype, **kwargs)
    818 nonzero_p = pt.as_tensor_variable(nonzero_p)
    819 weights = pt.stack([1 - nonzero_p, nonzero_p], axis=-1)
    820 comp_dists = [
    821     DiracDelta.dist(zero),
--> 822     Truncated.dist(nonzero_dist, lower=lower),
    823 ]
    825 if name is not None:
    826     return Mixture(name, weights, comp_dists, **kwargs)

File ~/anaconda3/envs/bambi-dev/lib/python3.11/site-packages/pymc/distributions/truncated.py:316, in Truncated.dist(cls, dist, lower, upper, max_n_steps, **kwargs)
    313 if lower is None and upper is None:
    314     raise ValueError("lower and upper cannot both be None")
--> 316 return super().dist([dist, lower, upper, max_n_steps], **kwargs)

File ~/anaconda3/envs/bambi-dev/lib/python3.11/site-packages/pymc/distributions/distribution.py:633, in Distribution.dist(cls, dist_params, shape, **kwargs)
    631 ndim_supp = getattr(cls.rv_type, "ndim_supp", None)
    632 if ndim_supp is None:
--> 633     ndim_supp = cls.rv_op(*dist_params, **kwargs).owner.op.ndim_supp
    634 create_size = find_size(shape=shape, size=size, ndim_supp=ndim_supp)
    635 rv_out = cls.rv_op(*dist_params, size=create_size, **kwargs)

File ~/anaconda3/envs/bambi-dev/lib/python3.11/site-packages/pymc/distributions/truncated.py:183, in TruncatedRV.rv_op(cls, dist, lower, upper, max_n_steps, size)
    179     return graph_inputs.index(rng)
    181 next_rngs = [next_rng for rng, next_rng in sorted(updates.items(), key=sort_updates)]
--> 183 return TruncatedRV(
    184     base_rv_op=dist.owner.op,
    185     inputs=graph_inputs,
    186     outputs=[truncated_rv, *next_rngs],
    187     ndim_supp=0,
    188     max_n_steps=max_n_steps,
    189 )(*graph_inputs)

File ~/anaconda3/envs/bambi-dev/lib/python3.11/site-packages/pymc/distributions/truncated.py:77, in TruncatedRV.__init__(self, base_rv_op, max_n_steps, *args, **kwargs)
     72 self.max_n_steps = max_n_steps
     73 self._print_name = (
     74     f"Truncated{self.base_rv_op._print_name[0]}",
     75     f"\\operatorname{{{self.base_rv_op._print_name[1]}}}",
     76 )
---> 77 super().__init__(*args, **kwargs)

File ~/anaconda3/envs/bambi-dev/lib/python3.11/site-packages/pymc/distributions/distribution.py:422, in SymbolicRandomVariable.__init__(self, *args, **kwargs)
    420 kwargs.setdefault("inline", True)
    421 kwargs.setdefault("strict", True)
--> 422 super().__init__(*args, **kwargs)

File ~/anaconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/compile/builders.py:423, in OpFromGraph.__init__(self, inputs, outputs, inline, lop_overrides, grad_overrides, rop_overrides, connection_pattern, strict, name, **kwargs)
    418 self.fgraph, self.shared_inputs, _, _ = construct_nominal_fgraph(
    419     inputs, outputs
    420 )
    422 if strict and self.shared_inputs:
--> 423     raise ValueError(
    424         "All variables needed to compute inner-graph must be provided as inputs under strict=True. "
    425         f"The inner-graph implicitly depends on the following shared variables {self.shared_inputs}"
    426     )
    428 self.kwargs = kwargs
    429 self.input_types = [inp.type for inp in inputs]

ValueError: All variables needed to compute inner-graph must be provided as inputs under strict=True. The inner-graph implicitly depends on the following shared variables [RandomGeneratorSharedVariable(<Generator(PCG64) at 0x71D9FE9BE420>)]

ValueError: All variables needed to compute inner-graph must be provided as inputs under strict=True. The inner-graph implicitly depends on the following shared variables [RandomGeneratorSharedVariable(<Generator(PCG64) at 0x71D9FE9BE420>)]

PyMC version information:

PyMC = 5.14.0
PyTensor = 2.20.0

Context for the issue:

I'm trying to finalize a large refactor in Bambi and some tests about HurdlePoisson didn't pass and that's how I found this.
The same thing happened with HurdleNegativeBinomial.

[tomicapretto]

I found a similar error message in a pretty big model I'm working with, and I think it's related to this same issue. I see the same error when I try to sample from the posterior predictive when I use pm.Truncated together with pm.CustomDist. See the following example:

import numpy as np
import pymc as pm
import pytensor.tensor as pt

from pymc.model.fgraph import clone_model

# simulate data
y_values = pm.draw(pm.Truncated.dist(pm.Exponential.dist(scale=[2, 4, 6]), upper=7), 200, random_seed=1234)
y_values = y_values.T.flatten()

groups = list("ABC")
groups_idx = np.repeat([0, 1, 2], 200)

assert len(y_values) == len(groups_idx)

coords = {
    "group": groups,
    "__obs__": np.arange(len(y_values))
}

# Works
with pm.Model(coords=coords) as model:
    groups_idx_data = pm.Data("groups_idx", groups_idx, dims="__obs__")
    
    b = pm.Normal("b", dims="group")
    scale = pm.Deterministic("scale", b[groups_idx_data], dims="__obs__")
    value_latent = pm.Exponential.dist(scale=scale)
    value = pm.Truncated("value", value_latent, upper=7, observed=y_values, dims="__obs__")

    idata = pm.sample(chains=2, random_seed=1234)


# Works
new_coords = {
    "__obs__": np.arange(3) + 100,
}

new_data = {
    "groups_idx": np.array([0, 1, 2])
}

with clone_model(model) as c_model:
    pm.set_data(new_data, coords=new_coords)
    predictions = pm.sample_posterior_predictive(
        idata, 
        var_names=["value"], 
        predictions=True,
        random_seed=1234,
    )

# Fails
new_coords = {
    "__obs__": np.arange(3) + 100,
}

new_data = {
    "groups_idx": np.array([0, 1, 2])
}

def f_exp(scale, size):
    return pm.Exponential.dist(scale=scale, size=size)

with clone_model(model) as c_model:
    pm.set_data(new_data, coords=new_coords)
    b = c_model["b"]
    groups_idx_data = c_model["groups_idx"]
    scale = pm.Deterministic("b_new", b[groups_idx_data], dims="__obs__")
    value_latent = pm.CustomDist.dist(scale, dist=f_exp)
    pm.Truncated("value_new", value_latent, upper=7, dims="__obs__")
    predictions = pm.sample_posterior_predictive(
        idata, 
        var_names=["value_new"], 
        predictions=True,
        random_seed=1234,
    )

ValueError: All variables needed to compute inner-graph must be provided as inputs under strict=True. The inner-graph implicitly depends on the following shared variables [RandomGeneratorSharedVariable(<Generator(PCG64) at 0x7415037912A0>), group, groups_idx]

[ricardoV94]

Is the problem gone with the bugfix?

[tomicapretto]

Nope, I installed from main and I still had the problem.

[ricardoV94]

Is the clone_model stuff needed to reproduce the problem?

[tomicapretto]

Nope, the following also fails with the same message

# Fails
new_coords = {
    "__obs__": np.arange(3) + 100,
}

new_data = {
    "groups_idx": np.array([0, 1, 2])
}

def f_exp(scale, size):
    return pm.Exponential.dist(scale=scale, size=size)

with model:
    pm.set_data(new_data, coords=new_coords)
    b = model["b"]
    groups_idx_data = model["groups_idx"]
    scale = pm.Deterministic("b_new", b[groups_idx_data], dims="__obs__")
    value_latent = pm.CustomDist.dist(scale, dist=f_exp)
    pm.Truncated("value_new", value_latent, upper=7, dims="__obs__")
    predictions = pm.sample_posterior_predictive(
        idata, 
        var_names=["value_new"], 
        predictions=True,
        random_seed=1234,
    )

[ricardoV94]

If you can get an even smaller example without the set_data / multiple models that's even better :)

[tomicapretto]

@ricardoV94 I'll try :)

[ricardoV94]

Here is a MWE:

import pymc as pm

def f_exp(scale, size):
    return pm.Exponential.dist(scale=scale, size=size)

with pm.Model() as model:
    b = pm.Normal("b", shape=(3,))
    value_latent = pm.CustomDist.dist(b[[0, 0, 1, 1, 2, 2]], dist=f_exp)
    pm.Truncated("value_new", value_latent, upper=7)