Not sure how to handle warning in pymc v4.0.0b2

I’ve recently given v4.0.0b2 a try and have been really enjoying the much faster GPU sampling!

I find however that all my models give the same error I’m not really sure what to do about as it looks like it’s referencing an abstraction level below pymc itself (aesara) so I’m not sure what’s the ‘correct’ pymc-way to deal with it:

site-packages/aesara/graph/ UserWarning: Supervisor is not added. Please build a FunctionGraphvia aesara.compile.function.types.std_graph()or add the Supervisor class manually.
sub_prof = optimizer.optimize(fgraph)

Is there a step I’m missing in how to port my pymc v3 models across? An example code snippet which works fine (but throws that warning) is:

with pm.Model() as model:
    ids_p = pm.MutableData("ids", ids)
    target = pm.MutableData("target", sample["target"])
    intercept = pm.Normal('intercept', mu=0, sd=1)
    sigma_p = pm.HalfCauchy('sigma_p', 0.5)
    p_offset = pm.Normal('p_offset', mu=0, sd=1, shape=n_ids)
    p = pm.Deterministic("p", p_offset * sigma_p)
    eps = pm.HalfCauchy("eps", 0.5)
    y = pm.Normal("y", mu=intercept+ p[ids_p], sd=eps, observed=target)
    trace = pm.sampling_jax.sample_numpyro_nuts(chain_method="vectorized")

I tried googling the message but didn’t find anything relevant - any help would be much appreciated!

1 Like

It’s not an error, just a UserWarning. You are not doing anything wrong. Good news is that it should be sorted out in the next beta release!