Defining Vectorizable Helper Function given Posterior Draws (Forecasting with BVAR)

Hello,

I am currently trying to vectorize a function to produce forecasts for a BVAR/BVARMA model. For reference, I have been following the notebook here to obtain a simple BVAR model: Bayesian Vector Autoregressive Models — PyMC example gallery

my attempt at creating a function to obtain n-step ahead forecasts is below:

#X_train is the dataframe of observations,renamed here just for transparency

def _forecast(intercept, lag_coefs, noise, forecast=10):
    
    len_data = len(X_train)
    new_draws = np.zeros((X_train.shape[0]+forecast, X_train.shape[1]))
    # Fill the new array with the observed data
    new_draws[:X_train] = X_train[:]
    for i in range(forecast):
       beta = calc_ar_step(lag_coefs, n_eqs, n_lags, df)
       mean = intercept + beta
       new_draws[len_data+i] = rng.normal(mean, noise)
        
    # Replace all observed data with nan, so they don't show when we plot it
    new_draws[:-forecast-1] = np.nan
    return new_draws



# Vectorized forecast function to handle multiple parameter draws
forecast = np.vectorize(
    _forecast,
    signature=("(v),(l,v,v),(v)->(o,v)"),
    excluded=("forecast",),
)
    

However, I encounter an inconsistent size for core dimension error; which leads me to believe that somewhere I have mistakenly summed over an axis or mis applied a transpose-which is my current debug path, but perhaps i am misunderstanding some of the xarray machinery

Have you seen/tried the statespace VARMAX implementation? It has helpers for forecasts and IRFs that automate these issues.

Regardless of how you fit the model I think the easiest way to handle these post-estimation tasks is to re-cast the model in statespace form and just iterate on x_t = T x_{t-1} + d + \varepsilon_t. That’s what I do in this blog post, which doesn’t use PyMC-Statespace. It’s much easier to use an auxiliary forecasting model than to try to wrestle with the xarrays. Here’s an example where I wrestled with the vectorized ufuncs but then switched to a forecasting model for an easier/happier life.

1 Like

Hey There! Thanks for your reply

I have not considered at this time to re-write as a state space, that’s an excellent point! My model is a dirichlet varmax technically, so its worth a shot just starting from the ground up and going from there! My main goal yesterday was to work on functions to help me model the additive log transform of the data + extract predictions

i’lll make an attempt today and edit this post if I make progress! thank you so much for your rec!

Hi Again!

Sorry for the debugging question; thus far my code compiles fine-but when i move to the prediction step however i’ve encountered an error i’m not super familiar with:

TypingError: Failed in nopython mode pipeline (step: nopython frontend)
- Resolution failure for literal arguments:
reshape() supports contiguous array only
- Resolution failure for non-literal arguments:
None

During: resolving callee type: BoundFunction(array.reshape for array(float64, 3d, A))
During: typing of call at C:\Users\broth\Forecasting_Scorecard\VAR_to_statespace.py (10)


File "VAR_to_statespace.py", line 10:
def make_statespace(coefs, intercepts):
    <source elided>
    T = np.zeros((n_states, n_states))
    T[:n_obs, :] = coefs.reshape(n_obs, n_states)
    ^

During: resolving callee type: type(CPUDispatcher(<function make_statespace at 0x00000259600191C0>))
During: typing of call at C:\Users\broth\Forecasting_Scorecard\VAR_to_statespace.py (79)

During: resolving callee type: type(CPUDispatcher(<function make_statespace at 0x00000259600191C0>))
During: typing of call at C:\Users\broth\Forecasting_Scorecard\VAR_to_statespace.py (79)

I haven’t used numba much at all-so I’m a bit at a lost for how i should perceive

I don’t think my data is an issue…my domain is R

Again, I don’t recommend you work with the draws directly. Instead, create the required statespace matrices using a new PyMC model and use pm.sample_posterior_predictive. Check the linked blog post for full explanation. This notebook might also clarify what you need to do.

I’ll reiterate that pymc_experimental.statepsace.BayesianVARMA handles all this for you, so I strongly recommend people use that.

Apologies; I misunderstood which link you were recommending; ill have a look at your new recommendations right now! I thought you were reccomending this approach VAR-Blog/Big VAR Blog Post.ipynb at main · jessegrabowski/VAR-Blog · GitHub, which which does not use the var package from experimental

I’d like to eventually take sums of categories and draw posteriors for them, so more flexibility is always nice (Here i’d be feeding the mean for two response additive log transformed categories and then passing them to a dirichlet distribution to output probabilities)

Hmmm, similar error out, i’m going through numba right now-apologies again I have not used this package before.

again I appreciate your help!

TypingError: Failed in nopython mode pipeline (step: nopython frontend)
Failed in nopython mode pipeline (step: nopython frontend)
Failed in nopython mode pipeline (step: nopython frontend)
No implementation of function Function(<function numba_funcify_Elemwise.<locals>.elemwise at 0x00000173031034C0>) found for signature:
 
elemwise(float64, readonly array(float64, 0d, C), array(float64, 0d, C))
 
There are 2 candidate implementations:
  - Of which 2 did not match due to:
  Overload in function 'numba_funcify_Elemwise.<locals>.ov_elemwise': File: pytensor\link\numba\dispatch\elemwise.py: Line 541.
    With argument(s): '(float64, readonly array(float64, 0d, C), array(float64, 0d, C))':
   Rejected as the implementation raised a specific error:
     TypingError: Failed in nopython mode pipeline (step: nopython frontend)
   No implementation of function Function(<intrinsic _vectorized>) found for signature:
    
_vectorized(type(CPUDispatcher(<function store_core_outputs at 0x00000173031267A0>)), Literal[str](gASVBgAAAAAAAAApKSmHlC4=
   ), Literal[str](gASVBAAAAAAAAAAphZQu
   ), Literal[str](gASVDQAAAAAAAACMB2Zsb2F0NjSUhZQu
   ), Literal[str](gASVCQAAAAAAAABLAEsAhpSFlC4=
   ), Tuple(), StarArgTuple(float64, readonly array(float64, 0d, C), array(float64, 0d, C)), UniTuple(Tuple() x 1), none)
    
   There are 2 candidate implementations:
         - Of which 1 did not match due to:
         Intrinsic in function '_vectorized': File: pytensor\link\numba\dispatch\vectorize_codegen.py: Line 74.
           With argument(s): '(type(CPUDispatcher(<function store_core_outputs at 0x00000173031267A0>)), Literal[str](gASVBgAAAAAAAAApKSmHlC4=
         ), Literal[str](gASVBAAAAAAAAAAphZQu
         ), Literal[str](gASVDQAAAAAAAACMB2Zsb2F0NjSUhZQu
         ), Literal[str](gASVCQAAAAAAAABLAEsAhpSFlC4=
         ), Tuple(), StarArgTuple(float64, readonly array(float64, 0d, C), array(float64, 0d, C)), UniTuple(Tuple() x 1), none)':
          Rejected as the implementation raised a specific error:
            TypingError: Vectorized inputs must be arrays.
     raised from C:\Users\broth\Anaconda3\envs\pymc_env\Lib\site-packages\pytensor\link\numba\dispatch\vectorize_codegen.py:130
         - Of which 1 did not match due to:
         Intrinsic in function '_vectorized': File: pytensor\link\numba\dispatch\vectorize_codegen.py: Line 74.
           With argument(s): '(type(CPUDispatcher(<function store_core_outputs at 0x00000173031267A0>)), unicode_type, unicode_type, unicode_type, unicode_type, Tuple(), StarArgTuple(float64, readonly array(float64, 0d, C), array(float64, 0d, C)), UniTuple(Tuple() x 1), none)':
          Rejected as the implementation raised a specific error:
            TypingError: input_bc_patterns must be literal.
     raised from C:\Users\broth\Anaconda3\envs\pymc_env\Lib\site-packages\pytensor\link\numba\dispatch\vectorize_codegen.py:100
   
   During: resolving callee type: Function(<intrinsic _vectorized>)
   During: typing of call at C:\Users\broth\Anaconda3\envs\pymc_env\Lib\site-packages\pytensor\link\numba\dispatch\elemwise.py (501)
   
   
   File "..\Anaconda3\envs\pymc_env\Lib\site-packages\pytensor\link\numba\dispatch\elemwise.py", line 501:
       def elemwise_wrapper(*inputs):
           return _vectorized(
           ^

  raised from C:\Users\broth\Anaconda3\envs\pymc_env\Lib\site-packages\numba\core\typeinfer.py:1091

During: resolving callee type: Function(<function numba_funcify_Elemwise.<locals>.elemwise at 0x00000173031034C0>)
During: typing of call at C:\Users\broth\AppData\Local\Temp\tmpu1e9orlz (79)


File "..\AppData\Local\Temp\tmpu1e9orlz", line 79:
def numba_funcified_fgraph(nominal_tensor_variable, nominal_tensor_variable_2, nominal_tensor_variable_1, nominal_tensor_variable_6, nominal_tensor_variable_3, nominal_tensor_variable_4, nominal_tensor_variable_7, nominal_tensor_variable_5):
    <source elided>
    # Composite{(i1 + log(i0) + i2)}(Det.0, 1.8378770664093453, dot.0)
    tensor_variable_38 = elemwise_4(tensor_variable_29, tensor_constant_9, tensor_variable_33)
    ^

During: resolving callee type: type(CPUDispatcher(<function numba_funcified_fgraph at 0x00000173028D98A0>))
During: typing of call at C:\Users\broth\AppData\Local\Temp\tmpzmw16avz (30)

During: resolving callee type: type(CPUDispatcher(<function numba_funcified_fgraph at 0x00000173028D98A0>))
During: typing of call at C:\Users\broth\AppData\Local\Temp\tmpzmw16avz (30)

During: resolving callee type: type(CPUDispatcher(<function numba_funcified_fgraph at 0x00000173028D98A0>))
During: typing of call at C:\Users\broth\AppData\Local\Temp\tmpzmw16avz (30)


File "..\AppData\Local\Temp\tmpzmw16avz", line 30:
def scan(n_steps, outer_in_1, outer_in_2, outer_in_3, outer_in_4, outer_in_5, outer_in_6, outer_in_7, outer_in_8, outer_in_9, outer_in_10, outer_in_11, outer_in_12, outer_in_13):
    <source elided>

        (inner_out_0, inner_out_1, inner_out_2, inner_out_3, inner_out_4, inner_out_5, inner_out_6) = scan_inner_func(np.asarray(outer_in_1[i]), np.asarray(outer_in_2[i]), np.asarray(outer_in_3[i]), np.asarray(outer_in_4_sitsot_storage[(i) % outer_in_4_len]), np.asarray(outer_in_5_sitsot_storage[(i) % outer_in_5_len]), outer_in_11, outer_in_12, outer_in_13)
        ^

During: resolving callee type: type(CPUDispatcher(<function scan at 0x00000173030CA520>))
During: typing of call at C:\Users\broth\AppData\Local\Temp\tmp0xzb18i5 (401)

During: resolving callee type: type(CPUDispatcher(<function scan at 0x00000173030CA520>))
During: typing of call at C:\Users\broth\AppData\Local\Temp\tmp0xzb18i5 (401)


File "..\AppData\Local\Temp\tmp0xzb18i5", line 401:
def numba_funcified_fgraph(_unconstrained_point, data):
    <source elided>
    # Scan{forward_kalman_pass, while_loop=False, inplace=all}(Shape_i{0}.0, Composite{...}.0, Composite{...}.1, data, SetSubtensor{:stop}.0, SetSubtensor{:stop}.0, Composite{...}.0, Composite{...}.0, Composite{...}.0, Composite{...}.0, Shape_i{0}.0, DropDims{axis=0}.0, Composite{(0.5 * (i0 + i1))}.0, DimShuffle{order=[2,1]}.0)
    tensor_variable_237, tensor_variable_238, tensor_variable_239, tensor_variable_240, tensor_variable_241, tensor_variable_242, loglike_obs = scan(tensor_variable_6, tensor_variable_7, tensor_variable_8, data, tensor_variable_222, tensor_variable_223, tensor_variable_124, tensor_variable_116, tensor_variable_108, tensor_variable_100, tensor_variable_6, tensor_variable_55, tensor_variable_225, tensor_variable_54)
    ^

During: resolving callee type: type(CPUDispatcher(<function numba_funcified_fgraph at 0x00000173064C4A40>))
During: typing of call at C:\Users\broth\Anaconda3\envs\pymc_env\Lib\site-packages\nutpie\compile_pymc.py (448)

During: resolving callee type: type(CPUDispatcher(<function numba_funcified_fgraph at 0x00000173064C4A40>))
During: typing of call at C:\Users\broth\Anaconda3\envs\pymc_env\Lib\site-packages\nutpie\compile_pymc.py (448)

During: resolving callee type: type(CPUDispatcher(<function numba_funcified_fgraph at 0x00000173064C4A40>))
During: typing of call at C:\Users\broth\Anaconda3\envs\pymc_env\Lib\site-packages\nutpie\compile_pymc.py (448)

During: resolving callee type: type(CPUDispatcher(<function numba_funcified_fgraph at 0x00000173064C4A40>))
During: typing of call at C:\Users\broth\Anaconda3\envs\pymc_env\Lib\site-packages\nutpie\compile_pymc.py (448)

During: resolving callee type: type(CPUDispatcher(<function numba_funcified_fgraph at 0x00000173064C4A40>))
During: typing of call at C:\Users\broth\Anaconda3\envs\pymc_env\Lib\site-packages\nutpie\compile_pymc.py (448)

During: resolving callee type: type(CPUDispatcher(<function numba_funcified_fgraph at 0x00000173064C4A40>))
During: typing of call at C:\Users\broth\Anaconda3\envs\pymc_env\Lib\site-packages\nutpie\compile_pymc.py (448)```

You’re going to have to share specific model code, it’s really not clear to me what you’re trying to do

No worries. I am attempting to create a dirichlet-arma model, as outlined here: https://arxiv.org/pdf/2303.17478v3. The details are not of super importance; basically you take compositional data; perform a log additive transform on the data, and then model the data as VAR/VARMAX etc etc.

Here’s my code: I’ve just attempted to pass my data into the varmax example directly; as I also am only considering normal priors on the lag coefficients for now

import numpy as np
import statsmodels.api as sm
import pandas as pd

import pymc as pm
import pytensor.tensor as pt
import arviz as az

import matplotlib.pyplot as plt
import sys

sys.path.append("..")
import pymc_experimental.statespace as pmss
import re
import math

import time
import datetime
from datetime import datetime as dt

#simulated dataframe of log additive data, data are draws from the set of real numbers
sim_data = np.array([[ 0.13524486, -1.1680408 , -0.4627132 , -1.07898902],
       [ 0.13524486, -1.1680408 , -0.4627132 , -1.07898902],
       [ 0.17348344, -1.13016065, -0.86557991, -0.84368122],
       [ 0.17348344, -1.13016065, -0.86557991, -0.84368122],
       [ 0.09825569, -1.24614773, -0.79229954, -0.90967549],
       [ 0.17348344, -1.13016065, -0.86557991, -0.84368122],
       [ 0.09825569, -1.24614773, -0.79229954, -0.90967549],
       [ 0.19480222, -1.08613163, -0.47031765, -1.16346483],
       [ 0.13983526, -1.0769267 , -0.58610174, -0.64219121],
       [ 0.18496357, -1.02961942, -0.44430494, -1.18302963],
       [-0.07265442, -1.67731906, -1.06149012, -1.43064501],
       [-0.17185026, -1.65658469, -1.11519642, -1.16666575],
       [-0.07265442, -1.67731906, -1.06149012, -1.43064501],
       [ 0.18259519, -1.28785429, -0.59768774, -0.71472478],
       [ 0.07878088, -1.18575916, -0.7561965 , -0.87397953],
       [-0.15600425, -1.48918878, -1.11541941, -1.23405026],
       [ 0.13524486, -1.1680408 , -0.4627132 , -1.07898902],
       [-0.07265442, -1.67731906, -1.06149012, -1.43064501],
       [ 0.37208589, -1.08731273, -0.48835277, -0.69314718],
       [ 0.21577998, -1.18867311, -0.59162176, -0.9315582 ],
       [ 0.13563163, -1.44615475, -0.69314718, -0.7927479 ],
       [ 0.05446526, -0.98032509, -0.7411812 , -0.93307221],
       [ 0.13983526, -1.0769267 , -0.58610174, -0.64219121],
       [ 0.37208589, -1.08731273, -0.48835277, -0.69314718],
       [-0.16154443, -1.53493539, -1.10749137, -1.21694796],
       [ 0.22608521, -1.04445893, -0.33975093, -1.16807289],
       [ 0.17348344, -1.13016065, -0.86557991, -0.84368122],
       [ 0.26383774, -1.10663158, -0.46299946, -1.1923263 ],
       [ 0.09825569, -1.24614773, -0.79229954, -0.90967549],
       [ 0.21577998, -1.18867311, -0.59162176, -0.9315582 ],
       [ 0.07878088, -1.18575916, -0.7561965 , -0.87397953],
       [-0.1268721 , -1.57442339, -1.20918309, -1.14417826],
       [ 0.22608521, -1.04445893, -0.33975093, -1.16807289],
       [ 0.22608521, -1.04445893, -0.33975093, -1.16807289],
       [-0.15600425, -1.48918878, -1.11541941, -1.23405026],
       [ 0.11401858, -1.21097526, -0.56486303, -0.69798981],
       [ 0.16717795, -1.07947295, -0.67242931, -0.85917705],
       [-0.17185026, -1.65658469, -1.11519642, -1.16666575],
       [ 0.11401858, -1.21097526, -0.56486303, -0.69798981],
       [-0.1268721 , -1.57442339, -1.20918309, -1.14417826],
       [-0.07985733, -1.89600826, -1.01438565, -1.26001949],
       [ 0.13563163, -1.44615475, -0.69314718, -0.7927479 ],
       [-0.16086151, -1.7266788 , -1.15322196, -1.16358474],
       [ 0.22608521, -1.04445893, -0.33975093, -1.16807289],
       [-0.16086151, -1.7266788 , -1.15322196, -1.16358474],
       [ 0.09825569, -1.24614773, -0.79229954, -0.90967549],
       [ 0.17348344, -1.13016065, -0.86557991, -0.84368122],
       [ 0.18259519, -1.28785429, -0.59768774, -0.71472478],
       [ 0.26383774, -1.10663158, -0.46299946, -1.1923263 ],
       [ 0.26383774, -1.10663158, -0.46299946, -1.1923263 ]])

sim_data = pd.DataFrame(sim_data,columns=['series 1','series 2','series 3','series 4'])

def create_train_test(data:pd.DataFrame, prop:float):
    N=data.shape[0]
   
    N=int(round(N*prop,0))
    Train = data[:N]
    Test = data[N:]
    return Train, Test

X_train, X_test = create_train_test(X, .85)                

n_lags = 2

bvar_mod = pmss.BayesianVARMAX(
    endog_names=X_train.columns,
    order=(2, 0),
    stationary_initialization=False,
    measurement_error=False,
    filter_type="standard",
    verbose=True,
)

bvar_mod.param_dims.values()

bvar_mod.coords

x0_dims, P0_dims, state_cov_dims, ar_dims = bvar_mod.param_dims.values()
coords = bvar_mod.coords

ar_dims

data = X_train
with pm.Model(coords=coords) as var_mod:
    x0 = pm.Normal("x0", dims=x0_dims)
    P0_diag = pm.Gamma("P0_diag", alpha=2, beta=1, size=data.shape[1] * 2, dims=P0_dims[0])
    P0 = pm.Deterministic("P0", pt.diag(P0_diag), dims=P0_dims)

    state_chol, _, _ = pm.LKJCholeskyCov(
        "state_chol", eta=1, n=bvar_mod.k_posdef, sd_dist=pm.Exponential.dist(lam=1)
    )

    ar_params = pm.Normal("ar_params", mu=0, sigma=1, dims=ar_dims)
    state_cov = pm.Deterministic("state_cov", state_chol @ state_chol.T, dims=state_cov_dims)

    bvar_mod.build_statespace_graph(data,)
    idata = pm.sample(nuts_sampler="nutpie", chains=8, cores= 1, draws=500)

When I run this, the following error is passed to me; i a currently looking through numba documentation to try to understand better

The following parameters should be assigned priors inside a PyMC model block: 
	x0 -- shape: (8,), constraints: None, dims: ('state',)
	P0 -- shape: (8, 8), constraints: Positive Semi-definite, dims: ('state', 'state_aux')
	ar_params -- shape: (8, 2, 8), constraints: None, dims: ('observed_state', 'ar_lag', 'observed_state_aux')
	state_cov -- shape: (4, 4), constraints: Positive Semi-definite, dims: ('shock', 'shock_aux')
C:\Users\broth\Anaconda3\envs\pymc_env\Lib\site-packages\pymc_experimental\statespace\utils\data_tools.py:97: UserWarning: No frequency was specific on the data's DateTimeIndex.
  warnings.warn(NO_FREQ_INFO_WARNING)
C:\Users\broth\Anaconda3\envs\pymc_env\Lib\site-packages\pytensor\link\numba\dispatch\basic.py:377: UserWarning: Numba will use object mode to run AdvancedSetSubtensor's perform method
  warnings.warn(
C:\Users\broth\Anaconda3\envs\pymc_env\Lib\site-packages\pytensor\link\numba\dispatch\basic.py:377: UserWarning: Numba will use object mode to run AdvancedSetSubtensor's perform method
  warnings.warn(
C:\Users\broth\Anaconda3\envs\pymc_env\Lib\site-packages\pytensor\link\numba\dispatch\basic.py:377: UserWarning: Numba will use object mode to run AdvancedSetSubtensor's perform method
  warnings.warn(
C:\Users\broth\Anaconda3\envs\pymc_env\Lib\site-packages\pytensor\link\numba\dispatch\basic.py:657: UserWarning: Numba will use object mode to allow the `compute_uv` argument to `numpy.linalg.svd`.
  warnings.warn(
C:\Users\broth\Anaconda3\envs\pymc_env\Lib\site-packages\pytensor\link\numba\dispatch\basic.py:657: UserWarning: Numba will use object mode to allow the `compute_uv` argument to `numpy.linalg.svd`.
  warnings.warn(
C:\Users\broth\Anaconda3\envs\pymc_env\Lib\site-packages\pytensor\link\numba\dispatch\basic.py:377: UserWarning: Numba will use object mode to run AdvancedSetSubtensor's perform method
  warnings.warn(
C:\Users\broth\Anaconda3\envs\pymc_env\Lib\site-packages\pytensor\link\numba\dispatch\basic.py:657: UserWarning: Numba will use object mode to allow the `compute_uv` argument to `numpy.linalg.svd`.
  warnings.warn(
C:\Users\broth\Anaconda3\envs\pymc_env\Lib\site-packages\pytensor\link\numba\dispatch\basic.py:377: UserWarning: Numba will use object mode to run AdvancedSubtensor's perform method
  warnings.warn(
C:\Users\broth\Anaconda3\envs\pymc_env\Lib\site-packages\pytensor\link\numba\dispatch\basic.py:377: UserWarning: Numba will use object mode to run AdvancedSubtensor's perform method
  warnings.warn(
C:\Users\broth\Anaconda3\envs\pymc_env\Lib\site-packages\pytensor\link\numba\dispatch\basic.py:377: UserWarning: Numba will use object mode to run AdvancedSetSubtensor's perform method
  warnings.warn(
C:\Users\broth\Anaconda3\envs\pymc_env\Lib\site-packages\pytensor\link\numba\dispatch\basic.py:377: UserWarning: Numba will use object mode to run AdvancedSetSubtensor's perform method
  warnings.warn(
Traceback (most recent call last):

  File ~\Anaconda3\envs\pymc_env\Lib\site-packages\spyder_kernels\py3compat.py:356 in compat_exec
    exec(code, globals, locals)

  File c:\users\broth\forecasting_scorecard\untitled3.py:198
    idata = pm.sample(nuts_sampler="nutpie", chains=8, cores= 1, draws=500)

  File ~\Anaconda3\envs\pymc_env\Lib\site-packages\pymc\sampling\mcmc.py:725 in sample
    return _sample_external_nuts(

  File ~\Anaconda3\envs\pymc_env\Lib\site-packages\pymc\sampling\mcmc.py:307 in _sample_external_nuts
    compiled_model = nutpie.compile_pymc_model(model)

  File ~\Anaconda3\envs\pymc_env\Lib\site-packages\nutpie\compile_pymc.py:184 in compile_pymc_model
    logp_numba = numba.cfunc(c_sig, **kwargs)(logp_numba_raw)

  File ~\Anaconda3\envs\pymc_env\Lib\site-packages\numba\core\decorators.py:279 in wrapper
    res.compile()

  File ~\Anaconda3\envs\pymc_env\Lib\site-packages\numba\core\compiler_lock.py:35 in _acquire_compile_lock
    return func(*args, **kwargs)

  File ~\Anaconda3\envs\pymc_env\Lib\site-packages\numba\core\ccallback.py:68 in compile
    cres = self._compile_uncached()

  File ~\Anaconda3\envs\pymc_env\Lib\site-packages\numba\core\ccallback.py:82 in _compile_uncached
    return self._compiler.compile(sig.args, sig.return_type)

  File ~\Anaconda3\envs\pymc_env\Lib\site-packages\numba\core\dispatcher.py:129 in compile
    raise retval

  File ~\Anaconda3\envs\pymc_env\Lib\site-packages\numba\core\dispatcher.py:139 in _compile_cached
    retval = self._compile_core(args, return_type)

  File ~\Anaconda3\envs\pymc_env\Lib\site-packages\numba\core\dispatcher.py:152 in _compile_core
    cres = compiler.compile_extra(self.targetdescr.typing_context,

  File ~\Anaconda3\envs\pymc_env\Lib\site-packages\numba\core\compiler.py:751 in compile_extra
    return pipeline.compile_extra(func)

  File ~\Anaconda3\envs\pymc_env\Lib\site-packages\numba\core\compiler.py:445 in compile_extra
    return self._compile_bytecode()

  File ~\Anaconda3\envs\pymc_env\Lib\site-packages\numba\core\compiler.py:513 in _compile_bytecode
    return self._compile_core()

  File ~\Anaconda3\envs\pymc_env\Lib\site-packages\numba\core\compiler.py:492 in _compile_core
    raise e

  File ~\Anaconda3\envs\pymc_env\Lib\site-packages\numba\core\compiler.py:479 in _compile_core
    pm.run(self.state)

  File ~\Anaconda3\envs\pymc_env\Lib\site-packages\numba\core\compiler_machinery.py:368 in run
    raise patched_exception

  File ~\Anaconda3\envs\pymc_env\Lib\site-packages\numba\core\compiler_machinery.py:356 in run
    self._runPass(idx, pass_inst, state)

  File ~\Anaconda3\envs\pymc_env\Lib\site-packages\numba\core\compiler_lock.py:35 in _acquire_compile_lock
    return func(*args, **kwargs)

  File ~\Anaconda3\envs\pymc_env\Lib\site-packages\numba\core\compiler_machinery.py:311 in _runPass
    mutated |= check(pss.run_pass, internal_state)

  File ~\Anaconda3\envs\pymc_env\Lib\site-packages\numba\core\compiler_machinery.py:273 in check
    mangled = func(compiler_state)

  File ~\Anaconda3\envs\pymc_env\Lib\site-packages\numba\core\typed_passes.py:112 in run_pass
    typemap, return_type, calltypes, errs = type_inference_stage(

  File ~\Anaconda3\envs\pymc_env\Lib\site-packages\numba\core\typed_passes.py:93 in type_inference_stage
    errs = infer.propagate(raise_errors=raise_errors)

  File ~\Anaconda3\envs\pymc_env\Lib\site-packages\numba\core\typeinfer.py:1091 in propagate
    raise errors[0]

TypingError: Failed in nopython mode pipeline (step: nopython frontend)
Failed in nopython mode pipeline (step: nopython frontend)
Failed in nopython mode pipeline (step: nopython frontend)
No implementation of function Function(<function numba_funcify_Elemwise.<locals>.elemwise at 0x0000017317C1F100>) found for signature:
 
elemwise(float64, readonly array(float64, 0d, C), array(float64, 0d, C))
 
There are 2 candidate implementations:
  - Of which 2 did not match due to:
  Overload in function 'numba_funcify_Elemwise.<locals>.ov_elemwise': File: pytensor\link\numba\dispatch\elemwise.py: Line 541.
    With argument(s): '(float64, readonly array(float64, 0d, C), array(float64, 0d, C))':
   Rejected as the implementation raised a specific error:
     TypingError: Failed in nopython mode pipeline (step: nopython frontend)
   No implementation of function Function(<intrinsic _vectorized>) found for signature:
    
_vectorized(type(CPUDispatcher(<function store_core_outputs at 0x0000017317C1DBC0>)), Literal[str](gASVBgAAAAAAAAApKSmHlC4=
   ), Literal[str](gASVBAAAAAAAAAAphZQu
   ), Literal[str](gASVDQAAAAAAAACMB2Zsb2F0NjSUhZQu
   ), Literal[str](gASVCQAAAAAAAABLAEsAhpSFlC4=
   ), Tuple(), StarArgTuple(float64, readonly array(float64, 0d, C), array(float64, 0d, C)), UniTuple(Tuple() x 1), none)
    
   There are 2 candidate implementations:
         - Of which 1 did not match due to:
         Intrinsic in function '_vectorized': File: pytensor\link\numba\dispatch\vectorize_codegen.py: Line 74.
           With argument(s): '(type(CPUDispatcher(<function store_core_outputs at 0x0000017317C1DBC0>)), Literal[str](gASVBgAAAAAAAAApKSmHlC4=
         ), Literal[str](gASVBAAAAAAAAAAphZQu
         ), Literal[str](gASVDQAAAAAAAACMB2Zsb2F0NjSUhZQu
         ), Literal[str](gASVCQAAAAAAAABLAEsAhpSFlC4=
         ), Tuple(), StarArgTuple(float64, readonly array(float64, 0d, C), array(float64, 0d, C)), UniTuple(Tuple() x 1), none)':
          Rejected as the implementation raised a specific error:
            TypingError: Vectorized inputs must be arrays.
     raised from C:\Users\broth\Anaconda3\envs\pymc_env\Lib\site-packages\pytensor\link\numba\dispatch\vectorize_codegen.py:130
         - Of which 1 did not match due to:
         Intrinsic in function '_vectorized': File: pytensor\link\numba\dispatch\vectorize_codegen.py: Line 74.
           With argument(s): '(type(CPUDispatcher(<function store_core_outputs at 0x0000017317C1DBC0>)), unicode_type, unicode_type, unicode_type, unicode_type, Tuple(), StarArgTuple(float64, readonly array(float64, 0d, C), array(float64, 0d, C)), UniTuple(Tuple() x 1), none)':
          Rejected as the implementation raised a specific error:
            TypingError: input_bc_patterns must be literal.
     raised from C:\Users\broth\Anaconda3\envs\pymc_env\Lib\site-packages\pytensor\link\numba\dispatch\vectorize_codegen.py:100
   
   During: resolving callee type: Function(<intrinsic _vectorized>)
   During: typing of call at C:\Users\broth\Anaconda3\envs\pymc_env\Lib\site-packages\pytensor\link\numba\dispatch\elemwise.py (501)
   
   
   File "..\Anaconda3\envs\pymc_env\Lib\site-packages\pytensor\link\numba\dispatch\elemwise.py", line 501:
       def elemwise_wrapper(*inputs):
           return _vectorized(
           ^

  raised from C:\Users\broth\Anaconda3\envs\pymc_env\Lib\site-packages\numba\core\typeinfer.py:1091

During: resolving callee type: Function(<function numba_funcify_Elemwise.<locals>.elemwise at 0x0000017317C1F100>)
During: typing of call at C:\Users\broth\AppData\Local\Temp\tmpv4jtugby (79)


File "..\AppData\Local\Temp\tmpv4jtugby", line 79:
def numba_funcified_fgraph(nominal_tensor_variable, nominal_tensor_variable_2, nominal_tensor_variable_1, nominal_tensor_variable_6, nominal_tensor_variable_3, nominal_tensor_variable_4, nominal_tensor_variable_7, nominal_tensor_variable_5):
    <source elided>
    # Composite{(i1 + log(i0) + i2)}(Det.0, 1.8378770664093453, dot.0)
    tensor_variable_38 = elemwise_4(tensor_variable_29, tensor_constant_9, tensor_variable_33)
    ^

During: resolving callee type: type(CPUDispatcher(<function numba_funcified_fgraph at 0x0000017317F50900>))
During: typing of call at C:\Users\broth\AppData\Local\Temp\tmpw513poex (30)

During: resolving callee type: type(CPUDispatcher(<function numba_funcified_fgraph at 0x0000017317F50900>))
During: typing of call at C:\Users\broth\AppData\Local\Temp\tmpw513poex (30)

During: resolving callee type: type(CPUDispatcher(<function numba_funcified_fgraph at 0x0000017317F50900>))
During: typing of call at C:\Users\broth\AppData\Local\Temp\tmpw513poex (30)


File "..\AppData\Local\Temp\tmpw513poex", line 30:
def scan(n_steps, outer_in_1, outer_in_2, outer_in_3, outer_in_4, outer_in_5, outer_in_6, outer_in_7, outer_in_8, outer_in_9, outer_in_10, outer_in_11, outer_in_12, outer_in_13):
    <source elided>

        (inner_out_0, inner_out_1, inner_out_2, inner_out_3, inner_out_4, inner_out_5, inner_out_6) = scan_inner_func(np.asarray(outer_in_1[i]), np.asarray(outer_in_2[i]), np.asarray(outer_in_3[i]), np.asarray(outer_in_4_sitsot_storage[(i) % outer_in_4_len]), np.asarray(outer_in_5_sitsot_storage[(i) % outer_in_5_len]), outer_in_11, outer_in_12, outer_in_13)
        ^

During: resolving callee type: type(CPUDispatcher(<function scan at 0x00000173178D7420>))
During: typing of call at C:\Users\broth\AppData\Local\Temp\tmpe8q8d1_l (401)

During: resolving callee type: type(CPUDispatcher(<function scan at 0x00000173178D7420>))
During: typing of call at C:\Users\broth\AppData\Local\Temp\tmpe8q8d1_l (401)


File "..\AppData\Local\Temp\tmpe8q8d1_l", line 401:
def numba_funcified_fgraph(_unconstrained_point, data):
    <source elided>
    # Scan{forward_kalman_pass, while_loop=False, inplace=all}(Shape_i{0}.0, Composite{...}.0, Composite{...}.1, data, SetSubtensor{:stop}.0, SetSubtensor{:stop}.0, Composite{...}.0, Composite{...}.0, Composite{...}.0, Composite{...}.0, Shape_i{0}.0, DropDims{axis=0}.0, Composite{(0.5 * (i0 + i1))}.0, DimShuffle{order=[2,1]}.0)
    tensor_variable_237, tensor_variable_238, tensor_variable_239, tensor_variable_240, tensor_variable_241, tensor_variable_242, loglike_obs = scan(tensor_variable_6, tensor_variable_7, tensor_variable_8, data, tensor_variable_222, tensor_variable_223, tensor_variable_124, tensor_variable_116, tensor_variable_108, tensor_variable_100, tensor_variable_6, tensor_variable_55, tensor_variable_225, tensor_variable_54)
    ^

During: resolving callee type: type(CPUDispatcher(<function numba_funcified_fgraph at 0x0000017317F689A0>))
During: typing of call at C:\Users\broth\Anaconda3\envs\pymc_env\Lib\site-packages\nutpie\compile_pymc.py (448)

During: resolving callee type: type(CPUDispatcher(<function numba_funcified_fgraph at 0x0000017317F689A0>))
During: typing of call at C:\Users\broth\Anaconda3\envs\pymc_env\Lib\site-packages\nutpie\compile_pymc.py (448)

During: resolving callee type: type(CPUDispatcher(<function numba_funcified_fgraph at 0x0000017317F689A0>))
During: typing of call at C:\Users\broth\Anaconda3\envs\pymc_env\Lib\site-packages\nutpie\compile_pymc.py (448)

During: resolving callee type: type(CPUDispatcher(<function numba_funcified_fgraph at 0x0000017317F689A0>))
During: typing of call at C:\Users\broth\Anaconda3\envs\pymc_env\Lib\site-packages\nutpie\compile_pymc.py (448)

During: resolving callee type: type(CPUDispatcher(<function numba_funcified_fgraph at 0x0000017317F689A0>))
During: typing of call at C:\Users\broth\Anaconda3\envs\pymc_env\Lib\site-packages\nutpie\compile_pymc.py (448)

During: resolving callee type: type(CPUDispatcher(<function numba_funcified_fgraph at 0x0000017317F689A0>))
During: typing of call at C:\Users\broth\Anaconda3\envs\pymc_env\Lib\site-packages\nutpie\compile_pymc.py (448)

No worries. I am attempting to create a dirichlet-arma model, as outlined here: https://arxiv.org/pdf/2303.17478v3. The details are not of super importance; basically you take compositional data; perform a log additive transform on the data, and then model the data as VAR/VARMAX etc etc.

I am currently on windows 10:
pymc 5.16.2 pypi_0 pypi
pymc-bart 0.5.14 pypi_0 pypi
pymc-experimental 0.1.1 pypi_0 pypi
numba 0.59.1 py311hf62ec03_0
numexpr 2.8.7 py311h1fcbade_0
numpy 1.25.2 py311h0b4df5a_0 conda-forge
numpydoc 1.7.0 py311haa95532_0
nutpie 0.9.1 py311h2282fe9_0 conda-forge

I have to use nutpie and resort to using a single core due to windows :frowning:
Here’s my code: I’ve just attempted to pass my data into the varmax example directly; as I also am only considering normal priors on the lag coefficients for now

import numpy as np
import statsmodels.api as sm
import pandas as pd

import pymc as pm
import pytensor.tensor as pt
import arviz as az

import matplotlib.pyplot as plt
import sys

sys.path.append("..")
import pymc_experimental.statespace as pmss
import re
import math

import time
import datetime
from datetime import datetime as dt

#simulated dataframe of log additive data, data are draws from the set of real numbers
sim_data = np.array([[ 0.13524486, -1.1680408 , -0.4627132 , -1.07898902],
       [ 0.13524486, -1.1680408 , -0.4627132 , -1.07898902],
       [ 0.17348344, -1.13016065, -0.86557991, -0.84368122],
       [ 0.17348344, -1.13016065, -0.86557991, -0.84368122],
       [ 0.09825569, -1.24614773, -0.79229954, -0.90967549],
       [ 0.17348344, -1.13016065, -0.86557991, -0.84368122],
       [ 0.09825569, -1.24614773, -0.79229954, -0.90967549],
       [ 0.19480222, -1.08613163, -0.47031765, -1.16346483],
       [ 0.13983526, -1.0769267 , -0.58610174, -0.64219121],
       [ 0.18496357, -1.02961942, -0.44430494, -1.18302963],
       [-0.07265442, -1.67731906, -1.06149012, -1.43064501],
       [-0.17185026, -1.65658469, -1.11519642, -1.16666575],
       [-0.07265442, -1.67731906, -1.06149012, -1.43064501],
       [ 0.18259519, -1.28785429, -0.59768774, -0.71472478],
       [ 0.07878088, -1.18575916, -0.7561965 , -0.87397953],
       [-0.15600425, -1.48918878, -1.11541941, -1.23405026],
       [ 0.13524486, -1.1680408 , -0.4627132 , -1.07898902],
       [-0.07265442, -1.67731906, -1.06149012, -1.43064501],
       [ 0.37208589, -1.08731273, -0.48835277, -0.69314718],
       [ 0.21577998, -1.18867311, -0.59162176, -0.9315582 ],
       [ 0.13563163, -1.44615475, -0.69314718, -0.7927479 ],
       [ 0.05446526, -0.98032509, -0.7411812 , -0.93307221],
       [ 0.13983526, -1.0769267 , -0.58610174, -0.64219121],
       [ 0.37208589, -1.08731273, -0.48835277, -0.69314718],
       [-0.16154443, -1.53493539, -1.10749137, -1.21694796],
       [ 0.22608521, -1.04445893, -0.33975093, -1.16807289],
       [ 0.17348344, -1.13016065, -0.86557991, -0.84368122],
       [ 0.26383774, -1.10663158, -0.46299946, -1.1923263 ],
       [ 0.09825569, -1.24614773, -0.79229954, -0.90967549],
       [ 0.21577998, -1.18867311, -0.59162176, -0.9315582 ],
       [ 0.07878088, -1.18575916, -0.7561965 , -0.87397953],
       [-0.1268721 , -1.57442339, -1.20918309, -1.14417826],
       [ 0.22608521, -1.04445893, -0.33975093, -1.16807289],
       [ 0.22608521, -1.04445893, -0.33975093, -1.16807289],
       [-0.15600425, -1.48918878, -1.11541941, -1.23405026],
       [ 0.11401858, -1.21097526, -0.56486303, -0.69798981],
       [ 0.16717795, -1.07947295, -0.67242931, -0.85917705],
       [-0.17185026, -1.65658469, -1.11519642, -1.16666575],
       [ 0.11401858, -1.21097526, -0.56486303, -0.69798981],
       [-0.1268721 , -1.57442339, -1.20918309, -1.14417826],
       [-0.07985733, -1.89600826, -1.01438565, -1.26001949],
       [ 0.13563163, -1.44615475, -0.69314718, -0.7927479 ],
       [-0.16086151, -1.7266788 , -1.15322196, -1.16358474],
       [ 0.22608521, -1.04445893, -0.33975093, -1.16807289],
       [-0.16086151, -1.7266788 , -1.15322196, -1.16358474],
       [ 0.09825569, -1.24614773, -0.79229954, -0.90967549],
       [ 0.17348344, -1.13016065, -0.86557991, -0.84368122],
       [ 0.18259519, -1.28785429, -0.59768774, -0.71472478],
       [ 0.26383774, -1.10663158, -0.46299946, -1.1923263 ],
       [ 0.26383774, -1.10663158, -0.46299946, -1.1923263 ]])

sim_data = pd.DataFrame(sim_data,columns=['series 1','series 2','series 3','series 4'])

def create_train_test(data:pd.DataFrame, prop:float):
    N=data.shape[0]
   
    N=int(round(N*prop,0))
    Train = data[:N]
    Test = data[N:]
    return Train, Test

X_train, X_test = create_train_test(X, .85)                

n_lags = 2

bvar_mod = pmss.BayesianVARMAX(
    endog_names=X_train.columns,
    order=(2, 0),
    stationary_initialization=False,
    measurement_error=False,
    filter_type="standard",
    verbose=True,
)

bvar_mod.param_dims.values()

bvar_mod.coords

x0_dims, P0_dims, state_cov_dims, ar_dims = bvar_mod.param_dims.values()
coords = bvar_mod.coords

ar_dims

data = X_train
with pm.Model(coords=coords) as var_mod:
    x0 = pm.Normal("x0", dims=x0_dims)
    P0_diag = pm.Gamma("P0_diag", alpha=2, beta=1, size=data.shape[1] * 2, dims=P0_dims[0])
    P0 = pm.Deterministic("P0", pt.diag(P0_diag), dims=P0_dims)

    state_chol, _, _ = pm.LKJCholeskyCov(
        "state_chol", eta=1, n=bvar_mod.k_posdef, sd_dist=pm.Exponential.dist(lam=1)
    )

    ar_params = pm.Normal("ar_params", mu=0, sigma=1, dims=ar_dims)
    state_cov = pm.Deterministic("state_cov", state_chol @ state_chol.T, dims=state_cov_dims)

    bvar_mod.build_statespace_graph(data,)
    idata = pm.sample(nuts_sampler="nutpie", chains=8, cores= 1, draws=500)

When I run this, the following error is passed to me; i a currently looking through numba documentation to try to understand better

The following parameters should be assigned priors inside a PyMC model block: 
	x0 -- shape: (8,), constraints: None, dims: ('state',)
	P0 -- shape: (8, 8), constraints: Positive Semi-definite, dims: ('state', 'state_aux')
	ar_params -- shape: (8, 2, 8), constraints: None, dims: ('observed_state', 'ar_lag', 'observed_state_aux')
	state_cov -- shape: (4, 4), constraints: Positive Semi-definite, dims: ('shock', 'shock_aux')
C:\Users\broth\Anaconda3\envs\pymc_env\Lib\site-packages\pymc_experimental\statespace\utils\data_tools.py:97: UserWarning: No frequency was specific on the data's DateTimeIndex.
  warnings.warn(NO_FREQ_INFO_WARNING)
C:\Users\broth\Anaconda3\envs\pymc_env\Lib\site-packages\pytensor\link\numba\dispatch\basic.py:377: UserWarning: Numba will use object mode to run AdvancedSetSubtensor's perform method
  warnings.warn(
C:\Users\broth\Anaconda3\envs\pymc_env\Lib\site-packages\pytensor\link\numba\dispatch\basic.py:377: UserWarning: Numba will use object mode to run AdvancedSetSubtensor's perform method
  warnings.warn(
C:\Users\broth\Anaconda3\envs\pymc_env\Lib\site-packages\pytensor\link\numba\dispatch\basic.py:377: UserWarning: Numba will use object mode to run AdvancedSetSubtensor's perform method
  warnings.warn(
C:\Users\broth\Anaconda3\envs\pymc_env\Lib\site-packages\pytensor\link\numba\dispatch\basic.py:657: UserWarning: Numba will use object mode to allow the `compute_uv` argument to `numpy.linalg.svd`.
  warnings.warn(
C:\Users\broth\Anaconda3\envs\pymc_env\Lib\site-packages\pytensor\link\numba\dispatch\basic.py:657: UserWarning: Numba will use object mode to allow the `compute_uv` argument to `numpy.linalg.svd`.
  warnings.warn(
C:\Users\broth\Anaconda3\envs\pymc_env\Lib\site-packages\pytensor\link\numba\dispatch\basic.py:377: UserWarning: Numba will use object mode to run AdvancedSetSubtensor's perform method
  warnings.warn(
C:\Users\broth\Anaconda3\envs\pymc_env\Lib\site-packages\pytensor\link\numba\dispatch\basic.py:657: UserWarning: Numba will use object mode to allow the `compute_uv` argument to `numpy.linalg.svd`.
  warnings.warn(
C:\Users\broth\Anaconda3\envs\pymc_env\Lib\site-packages\pytensor\link\numba\dispatch\basic.py:377: UserWarning: Numba will use object mode to run AdvancedSubtensor's perform method
  warnings.warn(
C:\Users\broth\Anaconda3\envs\pymc_env\Lib\site-packages\pytensor\link\numba\dispatch\basic.py:377: UserWarning: Numba will use object mode to run AdvancedSubtensor's perform method
  warnings.warn(
C:\Users\broth\Anaconda3\envs\pymc_env\Lib\site-packages\pytensor\link\numba\dispatch\basic.py:377: UserWarning: Numba will use object mode to run AdvancedSetSubtensor's perform method
  warnings.warn(
C:\Users\broth\Anaconda3\envs\pymc_env\Lib\site-packages\pytensor\link\numba\dispatch\basic.py:377: UserWarning: Numba will use object mode to run AdvancedSetSubtensor's perform method
  warnings.warn(
Traceback (most recent call last):

  File ~\Anaconda3\envs\pymc_env\Lib\site-packages\spyder_kernels\py3compat.py:356 in compat_exec
    exec(code, globals, locals)

  File c:\users\broth\forecasting_scorecard\untitled3.py:198
    idata = pm.sample(nuts_sampler="nutpie", chains=8, cores= 1, draws=500)

  File ~\Anaconda3\envs\pymc_env\Lib\site-packages\pymc\sampling\mcmc.py:725 in sample
    return _sample_external_nuts(

  File ~\Anaconda3\envs\pymc_env\Lib\site-packages\pymc\sampling\mcmc.py:307 in _sample_external_nuts
    compiled_model = nutpie.compile_pymc_model(model)

  File ~\Anaconda3\envs\pymc_env\Lib\site-packages\nutpie\compile_pymc.py:184 in compile_pymc_model
    logp_numba = numba.cfunc(c_sig, **kwargs)(logp_numba_raw)

  File ~\Anaconda3\envs\pymc_env\Lib\site-packages\numba\core\decorators.py:279 in wrapper
    res.compile()

  File ~\Anaconda3\envs\pymc_env\Lib\site-packages\numba\core\compiler_lock.py:35 in _acquire_compile_lock
    return func(*args, **kwargs)

  File ~\Anaconda3\envs\pymc_env\Lib\site-packages\numba\core\ccallback.py:68 in compile
    cres = self._compile_uncached()

  File ~\Anaconda3\envs\pymc_env\Lib\site-packages\numba\core\ccallback.py:82 in _compile_uncached
    return self._compiler.compile(sig.args, sig.return_type)

  File ~\Anaconda3\envs\pymc_env\Lib\site-packages\numba\core\dispatcher.py:129 in compile
    raise retval

  File ~\Anaconda3\envs\pymc_env\Lib\site-packages\numba\core\dispatcher.py:139 in _compile_cached
    retval = self._compile_core(args, return_type)

  File ~\Anaconda3\envs\pymc_env\Lib\site-packages\numba\core\dispatcher.py:152 in _compile_core
    cres = compiler.compile_extra(self.targetdescr.typing_context,

  File ~\Anaconda3\envs\pymc_env\Lib\site-packages\numba\core\compiler.py:751 in compile_extra
    return pipeline.compile_extra(func)

  File ~\Anaconda3\envs\pymc_env\Lib\site-packages\numba\core\compiler.py:445 in compile_extra
    return self._compile_bytecode()

  File ~\Anaconda3\envs\pymc_env\Lib\site-packages\numba\core\compiler.py:513 in _compile_bytecode
    return self._compile_core()

  File ~\Anaconda3\envs\pymc_env\Lib\site-packages\numba\core\compiler.py:492 in _compile_core
    raise e

  File ~\Anaconda3\envs\pymc_env\Lib\site-packages\numba\core\compiler.py:479 in _compile_core
    pm.run(self.state)

  File ~\Anaconda3\envs\pymc_env\Lib\site-packages\numba\core\compiler_machinery.py:368 in run
    raise patched_exception

  File ~\Anaconda3\envs\pymc_env\Lib\site-packages\numba\core\compiler_machinery.py:356 in run
    self._runPass(idx, pass_inst, state)

  File ~\Anaconda3\envs\pymc_env\Lib\site-packages\numba\core\compiler_lock.py:35 in _acquire_compile_lock
    return func(*args, **kwargs)

  File ~\Anaconda3\envs\pymc_env\Lib\site-packages\numba\core\compiler_machinery.py:311 in _runPass
    mutated |= check(pss.run_pass, internal_state)

  File ~\Anaconda3\envs\pymc_env\Lib\site-packages\numba\core\compiler_machinery.py:273 in check
    mangled = func(compiler_state)

  File ~\Anaconda3\envs\pymc_env\Lib\site-packages\numba\core\typed_passes.py:112 in run_pass
    typemap, return_type, calltypes, errs = type_inference_stage(

  File ~\Anaconda3\envs\pymc_env\Lib\site-packages\numba\core\typed_passes.py:93 in type_inference_stage
    errs = infer.propagate(raise_errors=raise_errors)

  File ~\Anaconda3\envs\pymc_env\Lib\site-packages\numba\core\typeinfer.py:1091 in propagate
    raise errors[0]

TypingError: Failed in nopython mode pipeline (step: nopython frontend)
Failed in nopython mode pipeline (step: nopython frontend)
Failed in nopython mode pipeline (step: nopython frontend)
No implementation of function Function(<function numba_funcify_Elemwise.<locals>.elemwise at 0x0000017317C1F100>) found for signature:
 
elemwise(float64, readonly array(float64, 0d, C), array(float64, 0d, C))
 
There are 2 candidate implementations:
  - Of which 2 did not match due to:
  Overload in function 'numba_funcify_Elemwise.<locals>.ov_elemwise': File: pytensor\link\numba\dispatch\elemwise.py: Line 541.
    With argument(s): '(float64, readonly array(float64, 0d, C), array(float64, 0d, C))':
   Rejected as the implementation raised a specific error:
     TypingError: Failed in nopython mode pipeline (step: nopython frontend)
   No implementation of function Function(<intrinsic _vectorized>) found for signature:
    
_vectorized(type(CPUDispatcher(<function store_core_outputs at 0x0000017317C1DBC0>)), Literal[str](gASVBgAAAAAAAAApKSmHlC4=
   ), Literal[str](gASVBAAAAAAAAAAphZQu
   ), Literal[str](gASVDQAAAAAAAACMB2Zsb2F0NjSUhZQu
   ), Literal[str](gASVCQAAAAAAAABLAEsAhpSFlC4=
   ), Tuple(), StarArgTuple(float64, readonly array(float64, 0d, C), array(float64, 0d, C)), UniTuple(Tuple() x 1), none)
    
   There are 2 candidate implementations:
         - Of which 1 did not match due to:
         Intrinsic in function '_vectorized': File: pytensor\link\numba\dispatch\vectorize_codegen.py: Line 74.
           With argument(s): '(type(CPUDispatcher(<function store_core_outputs at 0x0000017317C1DBC0>)), Literal[str](gASVBgAAAAAAAAApKSmHlC4=
         ), Literal[str](gASVBAAAAAAAAAAphZQu
         ), Literal[str](gASVDQAAAAAAAACMB2Zsb2F0NjSUhZQu
         ), Literal[str](gASVCQAAAAAAAABLAEsAhpSFlC4=
         ), Tuple(), StarArgTuple(float64, readonly array(float64, 0d, C), array(float64, 0d, C)), UniTuple(Tuple() x 1), none)':
          Rejected as the implementation raised a specific error:
            TypingError: Vectorized inputs must be arrays.
     raised from C:\Users\broth\Anaconda3\envs\pymc_env\Lib\site-packages\pytensor\link\numba\dispatch\vectorize_codegen.py:130
         - Of which 1 did not match due to:
         Intrinsic in function '_vectorized': File: pytensor\link\numba\dispatch\vectorize_codegen.py: Line 74.
           With argument(s): '(type(CPUDispatcher(<function store_core_outputs at 0x0000017317C1DBC0>)), unicode_type, unicode_type, unicode_type, unicode_type, Tuple(), StarArgTuple(float64, readonly array(float64, 0d, C), array(float64, 0d, C)), UniTuple(Tuple() x 1), none)':
          Rejected as the implementation raised a specific error:
            TypingError: input_bc_patterns must be literal.
     raised from C:\Users\broth\Anaconda3\envs\pymc_env\Lib\site-packages\pytensor\link\numba\dispatch\vectorize_codegen.py:100
   
   During: resolving callee type: Function(<intrinsic _vectorized>)
   During: typing of call at C:\Users\broth\Anaconda3\envs\pymc_env\Lib\site-packages\pytensor\link\numba\dispatch\elemwise.py (501)
   
   
   File "..\Anaconda3\envs\pymc_env\Lib\site-packages\pytensor\link\numba\dispatch\elemwise.py", line 501:
       def elemwise_wrapper(*inputs):
           return _vectorized(
           ^

  raised from C:\Users\broth\Anaconda3\envs\pymc_env\Lib\site-packages\numba\core\typeinfer.py:1091

During: resolving callee type: Function(<function numba_funcify_Elemwise.<locals>.elemwise at 0x0000017317C1F100>)
During: typing of call at C:\Users\broth\AppData\Local\Temp\tmpv4jtugby (79)


File "..\AppData\Local\Temp\tmpv4jtugby", line 79:
def numba_funcified_fgraph(nominal_tensor_variable, nominal_tensor_variable_2, nominal_tensor_variable_1, nominal_tensor_variable_6, nominal_tensor_variable_3, nominal_tensor_variable_4, nominal_tensor_variable_7, nominal_tensor_variable_5):
    <source elided>
    # Composite{(i1 + log(i0) + i2)}(Det.0, 1.8378770664093453, dot.0)
    tensor_variable_38 = elemwise_4(tensor_variable_29, tensor_constant_9, tensor_variable_33)
    ^

During: resolving callee type: type(CPUDispatcher(<function numba_funcified_fgraph at 0x0000017317F50900>))
During: typing of call at C:\Users\broth\AppData\Local\Temp\tmpw513poex (30)

During: resolving callee type: type(CPUDispatcher(<function numba_funcified_fgraph at 0x0000017317F50900>))
During: typing of call at C:\Users\broth\AppData\Local\Temp\tmpw513poex (30)

During: resolving callee type: type(CPUDispatcher(<function numba_funcified_fgraph at 0x0000017317F50900>))
During: typing of call at C:\Users\broth\AppData\Local\Temp\tmpw513poex (30)


File "..\AppData\Local\Temp\tmpw513poex", line 30:
def scan(n_steps, outer_in_1, outer_in_2, outer_in_3, outer_in_4, outer_in_5, outer_in_6, outer_in_7, outer_in_8, outer_in_9, outer_in_10, outer_in_11, outer_in_12, outer_in_13):
    <source elided>

        (inner_out_0, inner_out_1, inner_out_2, inner_out_3, inner_out_4, inner_out_5, inner_out_6) = scan_inner_func(np.asarray(outer_in_1[i]), np.asarray(outer_in_2[i]), np.asarray(outer_in_3[i]), np.asarray(outer_in_4_sitsot_storage[(i) % outer_in_4_len]), np.asarray(outer_in_5_sitsot_storage[(i) % outer_in_5_len]), outer_in_11, outer_in_12, outer_in_13)
        ^

During: resolving callee type: type(CPUDispatcher(<function scan at 0x00000173178D7420>))
During: typing of call at C:\Users\broth\AppData\Local\Temp\tmpe8q8d1_l (401)

During: resolving callee type: type(CPUDispatcher(<function scan at 0x00000173178D7420>))
During: typing of call at C:\Users\broth\AppData\Local\Temp\tmpe8q8d1_l (401)


File "..\AppData\Local\Temp\tmpe8q8d1_l", line 401:
def numba_funcified_fgraph(_unconstrained_point, data):
    <source elided>
    # Scan{forward_kalman_pass, while_loop=False, inplace=all}(Shape_i{0}.0, Composite{...}.0, Composite{...}.1, data, SetSubtensor{:stop}.0, SetSubtensor{:stop}.0, Composite{...}.0, Composite{...}.0, Composite{...}.0, Composite{...}.0, Shape_i{0}.0, DropDims{axis=0}.0, Composite{(0.5 * (i0 + i1))}.0, DimShuffle{order=[2,1]}.0)
    tensor_variable_237, tensor_variable_238, tensor_variable_239, tensor_variable_240, tensor_variable_241, tensor_variable_242, loglike_obs = scan(tensor_variable_6, tensor_variable_7, tensor_variable_8, data, tensor_variable_222, tensor_variable_223, tensor_variable_124, tensor_variable_116, tensor_variable_108, tensor_variable_100, tensor_variable_6, tensor_variable_55, tensor_variable_225, tensor_variable_54)
    ^

During: resolving callee type: type(CPUDispatcher(<function numba_funcified_fgraph at 0x0000017317F689A0>))
During: typing of call at C:\Users\broth\Anaconda3\envs\pymc_env\Lib\site-packages\nutpie\compile_pymc.py (448)

During: resolving callee type: type(CPUDispatcher(<function numba_funcified_fgraph at 0x0000017317F689A0>))
During: typing of call at C:\Users\broth\Anaconda3\envs\pymc_env\Lib\site-packages\nutpie\compile_pymc.py (448)

During: resolving callee type: type(CPUDispatcher(<function numba_funcified_fgraph at 0x0000017317F689A0>))
During: typing of call at C:\Users\broth\Anaconda3\envs\pymc_env\Lib\site-packages\nutpie\compile_pymc.py (448)

During: resolving callee type: type(CPUDispatcher(<function numba_funcified_fgraph at 0x0000017317F689A0>))
During: typing of call at C:\Users\broth\Anaconda3\envs\pymc_env\Lib\site-packages\nutpie\compile_pymc.py (448)

During: resolving callee type: type(CPUDispatcher(<function numba_funcified_fgraph at 0x0000017317F689A0>))
During: typing of call at C:\Users\broth\Anaconda3\envs\pymc_env\Lib\site-packages\nutpie\compile_pymc.py (448)

During: resolving callee type: type(CPUDispatcher(<function numba_funcified_fgraph at 0x0000017317F689A0>))
During: typing of call at C:\Users\broth\Anaconda3\envs\pymc_env\Lib\site-packages\nutpie\compile_pymc.py (448)

You may get a more readable error if you try with the default pymc nuts sampler. Do you?

Huh, that actually allows me to start drawing samples. Interesting! TY!

however, i am getting massive performance issue drops. with four equations and 2 lags (4 chains, 500 draws) , my approximate time to complete is 30 mins…it’s probably 5 when i define my model in this manner: Bayesian Vector Autoregressive Models — PyMC example gallery

The issue is the Numba backend does nor support all forms of advanced indexing that the default C backend does. You can try to use the JAX backend with numpyro nuts sampler and see how that goes

1 Like

I strongly recommend the numpyro sampler if you’re using the statespace models. You should get good speeds using that.

Numba support is still being worked on.

Hi again! Thank you for getting back to me. Apologies again for bearing with me; currently trying to cram some knowledge of jax in tonight.

numpyro does in fact speed things up a lot! However, I now run into this error when I sample from the conditional posterior after using your reccomended state space set up (copied directly from your code)

idata_post = bvar_mod.sample_conditional_posterior(idata)
/local_disk0/.ephemeral_nfs/cluster_libraries/python/lib/python3.10/site-packages/pytensor/link/jax/linker.py:28: UserWarning: The RandomType SharedVariables [RNG(<Generator(PCG64) at 0x7F80FABA3680>), RNG(<Generator(PCG64) at 0x7F8157418BA0>), RNG(<Generator(PCG64) at 0x7F8157419380>), RNG(<Generator(PCG64) at 0x7F815741AB20>), RNG(<Generator(PCG64) at 0x7F815741A340>), RNG(<Generator(PCG64) at 0x7F80FABA1A80>), RNG(<Generator(PCG64) at 0x7F80FABA34C0>), RNG(<Generator(PCG64) at 0x7F80FABA0900>), RNG(<Generator(PCG64) at 0x7F80FABA0AC0>), RNG(<Generator(PCG64) at 0x7F80FABA33E0>), RNG(<Generator(PCG64) at 0x7F8157419FC0>)] will not be used in the compiled JAX graph. Instead a copy will be used.
  warnings.warn(
Sampling: [filtered_posterior, filtered_posterior_observed, predicted_posterior, predicted_posterior_observed, smoothed_posterior, smoothed_posterior_observed]
TypeError: Shapes must be 1D sequences of concrete values of integer type, got (Traced<ShapedArray(int64[])>with<DynamicJaxprTrace(level=1/0)>, Traced<ShapedArray(int64[])>with<DynamicJaxprTrace(level=1/0)>, Traced<ShapedArray(int64[])>with<DynamicJaxprTrace(level=1/0)>).
If using `jit`, try using `static_argnums` or applying `jit` to smaller subfunctions.
The error occurred while tracing the function jax_funcified_fgraph at /tmp/tmpgvhrcj7b:1 for jit. This concrete value was not available in Python because it depends on the value of the argument observed_state.
The error occurred while tracing the function jax_funcified_fgraph at /tmp/tmpgvhrcj7b:1 for jit. This concrete value was not available in Python because it depends on the value of the argument ar_lag.
The error occurred while tracing the function jax_funcified_fgraph at /tmp/tmpgvhrcj7b:1 for jit. This concrete value was not available in Python because it depends on the value of the argument observed_state_aux.
Apply node that caused the error: Scan{scan_fn&scan_fn, while_loop=False, inplace=none}(Cast{int32}.0, Subtensor{start:stop:step}.0, Subtensor{start:stop:step}.0, Subtensor{start:stop:step}.0, Subtensor{start:stop:step}.0, RNG(<Generator(PCG64) at 0x7F80F1F209E0>), RNG(<Generator(PCG64) at 0x7F80F1F20AC0>), Cast{int32}.0, Cast{int32}.0)
Toposort index: 389
Inputs types: [TensorType(int32, shape=()), TensorType(float64, shape=(None, 8)), TensorType(float64, shape=(None, 8, 8)), TensorType(float64, shape=(None, 4)), TensorType(float64, shape=(None, 4, 4)), RandomGeneratorType, RandomGeneratorType, TensorType(int32, shape=()), TensorType(int32, shape=())]
Inputs shapes: [(23, 4), 'No shapes', (), 'No shapes', 'No shapes', (), (), (), 'No shapes', 'No shapes', 'No shapes', 'No shapes', 'No shapes', 'No shapes', 'No shapes', 'No shapes']
Inputs strides: [(8, 184), 'No strides', (), 'No strides', 'No strides', (), (), (), 'No strides', 'No strides', 'No strides', 'No strides', 'No strides', 'No strides', 'No strides', 'No strides']
Inputs values: ['not shown', {'bit_generator': 1, 'state': {'state': -1875307623921742080, 'inc': -2117585609497974999}, 'has_uint32': 0, 'uinteger': 0, 'jax_state': array([3858338214, 3750760192], dtype=uint32)}, array(8), {'bit_generator': 1, 'state': {'state': 5892849408795900207, 'inc': 3467606628746858139}, 'has_uint32': 0, 'uinteger': 0, 'jax_state': array([1372035920, 3458627887], dtype=uint32)}, {'bit_generator': 1, 'state': {'state': -4804107393757949562, 'inc': 7567433675770801199}, 'has_uint32': 0, 'uinteger': 0, 'jax_state': array([3176423879, 1413140870], dtype=uint32)}, array(4), array(2), array(4), {'bit_generator': 1, 'state': {'state': 6642182048432215167, 'inc': -6602606547338022137}, 'has_uint32': 0, 'uinteger': 0, 'jax_state': array([1546503521, 2588365951], dtype=uint32)}, {'bit_generator': 1, 'state': {'state': -4968201944813567915, 'inc': -6434350206830792963}, 'has_uint32': 0, 'uinteger': 0, 'jax_state': array([3138217639, 1660649557], dtype=uint32)}, {'bit_generator': 1, 'state': {'state': 5623932624895122884, 'inc': 7668392534115438363}, 'has_uint32': 0, 'uinteger': 0, 'jax_state': array([1309423852, 3952778692], dtype=uint32)}, {'bit_generator': 1, 'state': {'state': 196117885448786772, 'inc': 9192249282833447751}, 'has_uint32': 0, 'uinteger': 0, 'jax_state': array([  45662253, 2152108884], dtype=uint32)}, {'bit_generator': 1, 'state': {'state': 4985419079470314668, 'inc': -5619302220929038176}, 'has_uint32': 0, 'uinteger': 0, 'jax_state': array([1160758333,  675837100], dtype=uint32)}, {'bit_generator': 1, 'state': {'state': 5052423089649604090, 'inc': 7414153455853755774}, 'has_uint32': 0, 'uinteger': 0, 'jax_state': array([1176358919, 4186691066], dtype=uint32)}, {'bit_generator': 1, 'state': {'state': 2408003874362962264, 'inc': 126647561998871904}, 'has_uint32': 0, 'uinteger': 0, 'jax_state': array([ 560657091, 4247466328], dtype=uint32)}, {'bit_generator': 1, 'state': {'state': -2526832023300368508, 'inc': 5498351235899043001}, 'has_uint32': 0, 'uinteger': 0, 'jax_state': array([3706643369, 2618922884], dtype=uint32)}]
Outputs clients: [[output[0](filtered_posterior)], [output[1](filtered_posterior_observed)], [output[11](Scan{scan_fn&scan_fn, while_loop=False, inplace=none}.2)], [output[12](Scan{scan_fn&scan_fn, while_loop=False, inplace=none}.3)]]

HINT: Re-running with most PyTensor optimizations disabled could provide a back-trace showing when this node was created. This can be done by setting the PyTensor flag 'optimizer=fast_compile'. If that does not work, PyTensor optimizations can be disabled with 'optimizer=None'.
HINT: Use the PyTensor flag `exception_verbosity=high` for a debug print-out and storage map footprint of this Apply node.
File /local_disk0/.ephemeral_nfs/cluster_libraries/python/lib/python3.10/site-packages/pytensor/link/utils.py:196, in streamline.<locals>.streamline_default_f()
    193 for thunk, node, old_storage in zip(
    194     thunks, order, post_thunk_old_storage
    195 ):
--> 196     thunk()
    197     for old_s in old_storage:
File /local_disk0/.ephemeral_nfs/cluster_libraries/python/lib/python3.10/site-packages/jax/_src/core.py:1702, in canonicalize_shape(shape, context)
   1700 except TypeError:
   1701   pass
-> 1702 raise _invalid_shape_error(shape, context)

from the description it looks like it’s telling me i have some maligned vectors out there…but I’m going to be honest I don’t know where to start here

Oh it looks like the latest round of updates aren’t released yet. Can you try installing the latest version of experimental directly from github with pip install git+https://github.com/pymc-devs/pymc-experimental.git ?

Sure thing! Quickly cloning my environment so i can pip. I’m working out of conda, so if there is a recommendation on your end using conda I’ll be happy to do that too!

We don’t have pymc_experimental on conda-forge yet unfortunately

This is a really interesting approach to identifying VAR models in the prior through a change-of-variables.

Sarah Heaps. 2023. Enforcing Stationarity through the Prior in Vector Autoregressions. JCGS.

There’d probably be a way to define a combined prior plus constraint that would let you put this into PyMC.

Heaps 2023 is a really great resource. I show a thumbnail sketch of how to implement the prior in the VAR blog notebook linked above. I added this to the statespace roadmap as well. Help definitely wanted on this.