Positive Distributions as mu Causing Value Error in MvNormal

Hello, the following code produces an error:

with pm.Model() as model:

    mu = pm.LogNormal('mu',mu=1,sigma=1,shape=(2,2))
    var = pm.InverseGamma('var',alpha=1,beta=1,shape=2)
    cov = pm.Beta('cov',alpha=1,beta=1)

    cov_matrix = pm.math.stack([var[0],cov,cov, var[1]]).reshape((2,2))

    pm.MvNormal('mvn',
                mu=mu[0],
                cov=cov_matrix,
                observed=np.random.randn(100,2),
                )
    
    pm.sample(random_seed=0)

The error:

<<!! BUG IN FGRAPH.REPLACE OR A LISTENER !!>> <class ‘ValueError’> Shape of data ((4,)) does not match shape of type ((2,)) local_subtensor_lift
ERROR (pytensor.graph.rewriting.basic): Rewrite failure due to: local_subtensor_lift
ERROR (pytensor.graph.rewriting.basic): node: Subtensor{i}(Composite{…}.0, 0)
ERROR (pytensor.graph.rewriting.basic): TRACEBACK:
ERROR (pytensor.graph.rewriting.basic): Traceback (most recent call last):
File “c:\Users\hkinn\miniconda3\envs\pigs\Lib\site-packages\pytensor\graph\rewriting\basic.py”, line 1968, in process_node
fgraph.replace_all_validate_remove( # type: ignore
File “c:\Users\hkinn\miniconda3\envs\pigs\Lib\site-packages\pytensor\graph\features.py”, line 626, in replace_all_validate_remove
chk = fgraph.replace_all_validate(replacements, reason=reason, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File “c:\Users\hkinn\miniconda3\envs\pigs\Lib\site-packages\pytensor\graph\features.py”, line 571, in replace_all_validate
fgraph.replace(r, new_r, reason=reason, verbose=False, **kwargs)
File “c:\Users\hkinn\miniconda3\envs\pigs\Lib\site-packages\pytensor\graph\fg.py”, line 479, in replace
new_var = var.type.filter_variable(new_var, allow_convert=True)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File “c:\Users\hkinn\miniconda3\envs\pigs\Lib\site-packages\pytensor\tensor\type.py”, line 268, in filter_variable
other = self.constant_type(type=self, data=other)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File “c:\Users\hkinn\miniconda3\envs\pigs\Lib\site-packages\pytensor\tensor\variable.py”, line 1052, in init
raise ValueError(
ValueError: Shape of data ((4,)) does not match shape of type ((2,))

The sampler does continue past this error and seemingly samples correctly.

This error can be corrected by firstly changing the shape of mu (I have it as (2,2) in the previous example as that is what is required for my work), like so:

with pm.Model() as model:

    mu = pm.LogNormal('mu',mu=1,sigma=1,shape=2)
    var = pm.InverseGamma('var',alpha=1,beta=1,shape=2)
    cov = pm.Beta('cov',alpha=1,beta=1)

    cov_matrix = pm.math.stack([var[0],cov,cov, var[1]]).reshape((2,2))

    pm.MvNormal('mvn',
                mu=mu,
                cov=cov_matrix,
                observed=np.random.randn(100,2),
                )
    
    pm.sample(random_seed=0)

The error can also be fixed by changing the LogNormal to a Normal (with shape still (2,2)):

with pm.Model() as model:

    mu = pm.Normal('mu',mu=1,sigma=1,shape=(2,2))
    var = pm.InverseGamma('var',alpha=1,beta=1,shape=2)
    cov = pm.Beta('cov',alpha=1,beta=1)

    cov_matrix = pm.math.stack([var[0],cov,cov, var[1]]).reshape((2,2))

    pm.MvNormal('mvn',
                mu=mu[0],
                cov=cov_matrix,
                observed=np.random.randn(100,2),
                )
    
    pm.sample(random_seed=0)

I assume it is something to do with the covariance matrix, if you replace it with a constant identity matrix the error also disappears:

with pm.Model() as model:

    mu = pm.LogNormal('mu',mu=1,sigma=1,shape=(2,2))
    pm.MvNormal('mvn',
                mu=mu[0],
                cov=np.eye(2),
                observed=np.random.randn(100,2),
                )
    
    pm.sample(random_seed=0)

You can reproduce the error with other positive distributions like Gamma and InverseGamma.

Since the error does not stop the sampling it is not a massive issue, but it would be nice to know if I am doing something wrong.

In my actual, more complex model, I am sometimes able to fix the error by using .reshape on some variables, but this seems to be without any logic and I just guess. In this simplified example I could not manage to fix it using .reshape.

Thanks in advance for the help.

Edit: it turns out it is a problem since the sampling with this error is very slow.

I don’t see the error locally. Sampling took ~1 minute for me, although there are divergences (my guess is that the LogNormal mu prior is a bit at odds with negative observed values).

Are you using the last version of PyMC? If you are not, can you try updating?

Hi, that environment was running pymc 5.10.0. I created a fresh environment with 5.10.3, and I no longer have the bugs. No worries about the divergences, I just created the example quickly to recreate the bug.

Unclear if the original environment was somehow corrupted or if it was just the update that fixed the problem.

Should have just tried that in the in the first place, thanks a lot for the help and sorry for the waste of time.

1 Like

We have fixed some similar-sounding bugs recently, so I think it was just that.

I still get the error on 5.10.2, so it must have been something recent. @ricardoV94 is always fixing things!