Sampling_jax and Timeseries AR on pymc

[luiswilbert]

Dear colleagues,

Hope you guys are okay.

I'm creating a forecast model which will include trend + autoregressive error (latent) + uncorrelated error. The trend will be a custom weibull function. The data is a matrix with product in rows and month as columns.

I have created the model using pymc3. However, the sampling is slow.
Model is below. Therefore, I have tried to use the sampling_jax to speed up the process.

I cannot find the sampling_jax in the pymc3 library. Then, I have installed the library pymc to be able to use the sampling_jax. When I run the model, there is a error on the autoregressive model pm.AR (TypeError: dist() missing 1 required positional argument: 'dist_params'). Please, see the model and the error below.

I'm using the follow link to install the libraries: !pip install git+https://github.com/pymc-devs/pymc.git

My questions are:

• Is there a way to find the sampling_jax in the pymc3 library? I'm using the version 3.11. When I tried to use the sampling_jax, I have found the error: module 'pymc3' has no attribute 'sampling_jax'
• How Can I fix this problem? there is a error on the autoregressive model pm.AR (TypeError: dist() missing 1 required positional argument: 'dist_params').
• Is there a way to avoid the for loop and work as a matrix?
• Does the AR model support to work with matrix/vector?
• Has the timeseries package (AR) already been rebuilt on pymc4?

Thank you in advance

Model import

import pandas as pd
import numpy as np
import theano.tensor as tt
import xarray
import pymc as pm
import pymc.sampling_jax

Data as matrix (rows -> products | columns -> months)

df=np.array([[0.08, 0.14, 0.18, 0.06, 0.45, 4.05, 2.69, 2.63, 1.24, 1.71],
[0.95, 0.31, 0.37, 0.18, 0.17, 0.09, 0.11, 0.04, 0. , 0. ],
[0.07, 0.02, 0.04, 1.13, 0.44, 1.89, 0.71, 0.09, 0.03, 0.02],
[0.32, 0.24, 0.26, 0.35, 0.29, 0.18, 0.15, 0.02, 0.03, 0.02]])

Custom Weibull Function -----------------------------------------------

def weibullf(m,b,c,x): #custom weibull function
return(m*(((x/c)**(bc))np.exp(-bx+bc)))

rows, cols = df.shape
t=np.arange(1,cols+1)

Model -----------------------------------------------

with pm.Model() as model:

###Custom Weibull Parameters -----------------------------------------------

m = pm.Beta('m',1.4,4,shape=(rows,1))       #level
c = pm.Uniform('c',1,36,shape=(rows,1))     #center
b = pm.Uniform('b',0,3,shape=(rows,1))      #scale

###Autoregressive error Parameters -----------------------------------------------

theta1 = pm.Normal("theta1", 0, 1.0, shape=(rows,1))
sigma = pm.HalfNormal("tau", 0.05,shape=(rows,1))

###Likelihood error -----------------------------------------------

error =  pm.HalfNormal("error", 0.05)

###For loop to create the variable for each product (row) -----------------------------------------------

for i in range(rows): #for loop for each products (rows)
   auto_error=pm.AR("AR_%d" % i,  rho=[theta1[i]], sd=sigma[i],shape=(len(df[i, :]))) #autoregressive error (latent)
   trend = weibullf(m[i],b[i],c[i],t) #custom weibull function                                                 
   mu=trend+auto_error   #trend + autoregressive error
   pm.Normal( "obs_%d" % i,mu=multiprocessing, sigma=error,observed=df[i])  #likelihood

###Sampling AJAX -----------------------------------------------
trace = pm.sampling_jax.sample_numpyro_nuts(target_accept=0.95, random_seed=100)

Error

TypeError Traceback (most recent call last)
in ()
21
22 for i in range(rows):
---> 23 auto_error=pm.AR("AR_%d" % i, rho=[theta1[i]], sd=sigma[i],shape=(len(df[i, :]))) #autoregressive error (latent)
24
25 trend = weibullf(m[i],b[i],c[i],t) #custom weibull function

/usr/local/lib/python3.7/dist-packages/pymc/distributions/distribution.py in new(cls, name, rng, dims, initval, observed, total_size, transform, *args, **kwargs)
225 # Create the RV without dims information, because that's not something tracked at the Aesara level.
226 # If necessary we'll later replicate to a different size implied by already known dims.
--> 227 rv_out = cls.dist(*args, rng=rng, **kwargs)
228 ndim_actual = rv_out.ndim
229 resize_shape = None

TypeError: dist() missing 1 required positional argument: 'dist_params'

[ricardoV94]

@luiswilbert timeseries distributions have not been refactored in the latest version yet. There is an open issue here to track it: #4642