Predicting from regression model that depends on T

Hello,

I’m trying to generate predictions from the following model:

def BQR(tau=0.5, a=5, b=.04, sigma_beta_a=3, sigma_beta_b=1, X=None, y=None):
    T, K = X.shape
    
    # Deterministic
    theta            = (1-2*tau)/(tau*(1-tau))
    tau_star_squared = 2/(tau*(1-tau))
    
    # Non-Beta Priors
    sigma = numpyro.sample('sigma', dist.InverseGamma(a,1/b))
    z     = numpyro.sample('z', dist.Exponential((1/sigma)*jnp.ones(T)))
    
    # Beta Priors
    sigma_beta     = numpyro.sample('sigma_beta', dist.InverseGamma(sigma_beta_a,sigma_beta_b))
    unscaled_betas = numpyro.sample("unscaled_betas", dist.Normal(0., jnp.ones(K)))
    beta           = numpyro.deterministic("beta",  sigma_beta * unscaled_betas)
    beta0          = numpyro.sample('beta0', dist.Normal(0., 1.))

    # Likelihood
    mean_function   = beta0+jnp.matmul(X,beta)+theta*z
    stddev_function = jnp.sqrt(tau_star_squared*sigma*z)
    numpyro.sample("y",dist.Normal(mean_function,stddev_function),obs=y)

I use the following code (where model is BQR from above):

mcmc_obj = MCMC(NUTS(BQR,
              max_tree_depth=20,
              target_accept_prob=0.99),
          num_warmup=5000,
          num_samples=10000,
          num_chains=1,
          progress_bar=False)
def forecast_loop_w_predictive(X,y, mcmc_obj, model, fTs):  
        
    for f in fTs:
        t, quantile = f
        # Training Data up to and including t
        y_train, X_train = y.copy().loc[:t].to_numpy(), X.copy().loc[:t].to_numpy()
        # Testing Data for t+1
        y_test, X_test   = y.copy().loc[t:].iloc[1]   , X.copy().loc[t:].iloc[[1]].to_numpy()
        print("Size of X_train =",X_train.shape)

        # Run model
        mcmc_obj.run(random.PRNGKey(0),tau=quantile, X=X_train,y=y_train)

        # Forecast
        predictive     = Predictive(model, posterior_samples=mcmc_obj.get_samples())
        y_pred_samples = predictive(random.PRNGKey(0), X=X_test)["y"]

        return y_pred_samples

where t is a date, X_train is T x K, y_test is a single value, and X_test is 1 x K, where K is the number of independent variables and T is the number of time periods. I expect y_pred_samples to be num_samples x 1 but instead it is num_samples x T.

I think the issue has to do with the fact that BQR depends on T, and even though X_test is 1 x K, it still “remembers” T from X_train.

My question is, how can I modify my BQR model so that when calling Predictive I get back a num_samples x 1 array of predictions for y.

Thank you in advance.

Your latent z has shape T. MCMC will give you a collection of (num_samples, T) z. What is your expected shape for it when you make prediction?