Problem with LKJCholeskyCov in BVAR models

Hi, I’m trying to create implementation of BVAR model, but I see that smth breaks when I’m trying to use pm.LKJCholeskyCov in my implementation. Could smn please point out how to traceback such errors?

import pymc 
from importlib import reload
reload(pymc)

pm = pymc

def create_my_model(data, n_train: int, n_lags: int, priors, sampler_params, mv_norm = True):

    if n_train>0:
        df = data.iloc[:n_train].copy(deep=True)
    else:
        df = data.copy(deep=True)
            

    n_obs = df.shape[0]    
    n_eq = df.shape[1] # Number of equations equal to number of columns

    coords = {
        "lags": np.arange(n_lags) + 1,
        "equations": df.columns.tolist(),
        "cross_vars": df.columns.tolist(),
        "time": [i for i in df.index[n_lags:]],
    }


    with pm.Model(coords=coords) as model:

        # N_eq equations, each with N_lags and N_eq variables (to which lags will be applied)        
        lag_coefs = pm.Normal(
            "lag_coefs",
            mu=priors["lag_coefs"]["mu"],
            sigma=priors["lag_coefs"]["sigma"],
            dims=["equations", "lags", "cross_vars"],
        )
    
        alpha = pm.Normal(
            "alpha", mu=priors["alpha"]["mu"], sigma=priors["alpha"]["sigma"], dims=("equations",)
        )

        data_obs = pm.Data(f"data_obs", 
                           df.iloc[n_lags:].values,  
                           mutable=True)
        
        data_lags = {}
        n_obs = df.shape[0]
        for i in range(n_lags):
            dl = pm.Data(f"data_lag{i+1}", 
                         df.iloc[n_lags-i-1:n_obs-i-1].values,  
                         mutable=True)

            print(dl.shape.eval())
            data_lags[i] = dl
        
        # Create VAR equations
        var = []
        # returns tensor with dimension [n_train_obs - n_lags]
        for j in range(n_eq):
            ar = pm.math.sum(
                [
                pm.math.sum(lag_coefs[j, i] * data_lags[i], axis=-1) 
                    for i in range(n_lags)
                ], axis=0)
            print ('ar_eq:', ar.shape.eval())
            var.append(ar)
    
        # beta multiplied by X
        # shape = [n_train_obs - n_lags, n_eq]
        betaX = pm.math.stack(var, axis=-1)
    
        print ('betax:', betaX.shape.eval())

        mean = alpha + betaX

        mean = pm.Deterministic(
            "alpha+betaX",
            mean
        )        
        print ('mean:', mean.shape.eval())
    
        if mv_norm:
            # https://www.pymc.io/projects/docs/en/stable/api/distributions/generated/pymc.LKJCholeskyCov.html
            chol, _, _ = pm.LKJCholeskyCov(
                "noise_chol",
                eta=priors["noise_chol"]["eta"],
                n=n_eq,
                sd_dist=pm.HalfNormal.dist(sigma=priors["noise_chol"]["sigma"],shape=n_eq), 
            )
            
            # print ("noise_chol", chol.shape.eval())

            # obs = pm.MvNormal(
            #     "obs", mu=mean, chol=chol, observed=data_obs)

            vals_raw = pm.Normal("vals_raw", mu=mean, sigma=1)
            vals = pm.math.dot(vals_raw, chol)
            obs = pm.Normal("obs", mu=vals, sigma=0.0001, observed=data_obs)       
            print ('obs', obs.shape.eval())
        
        else:
            ## This is an alternative likelihood that can recover sensible estimates of the coefficients
            ## But lacks the multivariate correlation between the timeseries.
            sigma = pm.HalfNormal("noise", sigma=priors["noise"]["sigma"], dims=["equations"])
                
            print ('sigma:', sigma.shape.eval())
            print ('data_obs:', data_obs.shape.eval())        
            obs = pm.Normal(
                "obs", mu=mean, sigma=sigma, observed=data_obs)

        idata = None
        idata = pm.sample(**sampler_params)     
        
    return model, idata


n_lags = 3
n_train = 100
mv_norm = True

priors = {
    "lag_coefs": {"mu": 0.3, "sigma": 1.0},
    "alpha": {"mu": 15.0, "sigma": 5.0},
    "noise_chol": {"eta": 1.0, "sigma": 1.0},
    "noise": {"sigma": 1.0},

}


sampler_params = {'chains': 2, 'draws': 200, 'tune': 200}


df_full = pd.DataFrame(np.random.random(size = [500, 4]))

model, idata = create_my_model(df_full, 
                               n_train = n_train, 
                               n_lags = n_lags, 
                               priors = priors, sampler_params = sampler_params, mv_norm = True) 

Code fails with ‘blockwise alloc error’ which doesn’t make any sence to me

ERROR (pytensor.graph.rewriting.basic): Rewrite failure due to: local_blockwise_alloc
ERROR (pytensor.graph.rewriting.basic): node: Blockwise{Tri{dtype='float64'}, (),(),()->(o00,o01)}(Alloc.0, Alloc.0, [0])
ERROR (pytensor.graph.rewriting.basic): TRACEBACK:
ERROR (pytensor.graph.rewriting.basic): Traceback (most recent call last):

  File "/Users/ivanpetrov/.pyenv/versions/3.10.13/envs/deep_env310/lib/python3.10/site-packages/pytensor/graph/rewriting/basic.py", line 1922, in process_node
    replacements = node_rewriter.transform(fgraph, node)
  File "/Users/ivanpetrov/.pyenv/versions/3.10.13/envs/deep_env310/lib/python3.10/site-packages/pytensor/graph/rewriting/basic.py", line 1082, in transform
    return self.fn(fgraph, node)
  File "/Users/ivanpetrov/.pyenv/versions/3.10.13/envs/deep_env310/lib/python3.10/site-packages/pytensor/tensor/rewriting/blockwise.py", line 199, in local_blockwise_alloc
    assert new_outs[0].type.broadcastable == old_out_type.broadcastable
AssertionError

Name: pymc
Version: 5.11.0

Name: pytensor
Version: 2.18.6

Python 3.10

I played around with it a bit and I think it might be a bug related to MutableData. Can you change all of your pm.Data(..., mutable=True) to pm.ConstantData(..) and check if the error persists?

Hi both,

I am also having this problem. I have a pretty simple use case - I want to estimate the parameters of a multivariate normal distribution based on some data, and then sample a posterior predictive distribution for larger numbers of data than I observed (and some other stuff, but this is the crux of it) by using pm.set_data

Heres a simple example that is giving me the same error:

import numpy as np
import pymc as pm

test_data = np.random.normal(50, 10, size=(400, 3))

with pm.Model() as test:

    # Set data as in doctstring
    observed_data = pm.Data('observed_data', test_data)

    # Prior for means
    μ = pm.Normal('μ', mu=[50, 50, 50], sigma=10, shape=3)

    # LKJ
    chol, corr, sigmas = pm.LKJCholeskyCov('ρ', n=3, eta=1, sd_dist=pm.HalfNormal.dist(3))

    # Likelihood
    obs = pm.MvNormal('obs', mu=μ, chol=chol, observed=observed_data, shape=observed_data.shape)

    idata = pm.sample()

On running this I get the following:

ERROR (pytensor.graph.rewriting.basic): Rewrite failure due to: local_blockwise_alloc
ERROR (pytensor.graph.rewriting.basic): node: Blockwise{Tri{dtype='float64'}, (),(),()->(o00,o01)}(Alloc.0, Alloc.0, [0])
ERROR (pytensor.graph.rewriting.basic): TRACEBACK:
ERROR (pytensor.graph.rewriting.basic): Traceback (most recent call last):
  File "/opt/miniconda3/envs/py11/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1919, in process_node
    replacements = node_rewriter.transform(fgraph, node)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/miniconda3/envs/py11/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/opt/miniconda3/envs/py11/lib/python3.11/site-packages/pytensor/tensor/rewriting/blockwise.py", line 191, in local_blockwise_alloc
    assert new_outs[0].type.broadcastable == old_out_type.broadcastable
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
AssertionError

This feels like it should work but I am probably missing something. I am on PyMC 5.13.1. Tagging @ricardoV94, unsure if its related to this thread? Any help appreciated, I am stuck!