I am trying to dynamically slice my hierarchical model to allow for new predictions on data of a different shape. The issue I am running into is I can’t find a way to slice my 2D tensor dynamically in the model stack to extract the necessary bits after the dot product. The below code triggers the same error I am running into that there are no arrays when doing advanced indexing.
After I compute my dot product, I only want to retain (or slice) out each element given the index in the groups.
I have tried to set the shape/pt.arange()
component both as a pm.Deterministic
and pm.MutabeData
structure to no avail.
import numpy as np
import pymc as pm
import pytensor.tensor as pt
np.random.seed(123)
n = 5
variables = 3
groups = 2
weights = np.random.randn(variables, groups)
x = np.random.randn(n, variables)
group_index = np.random.choice(groups, n)
variable_names = ['var'+str(i) for i in range(variables)]
xw = np.dot(x,weights)
xw = xw[np.arange(xw.shape[0]),group_index]
y = xw + np.random.randn(n)/50
with pm.Model() as model:
model.add_coord('obs', np.arange(n), mutable=True)
model.add_coord('var_names', variable_names, mutable=False)
model.add_coord('group_values', values=[0,1], mutable=False)
data = pm.MutableData("data", x, dims=('obs','var_names'))
sigma_variables = pm.HalfNormal("sigma_variables", sigma=1, dims = 'group_values')
mu_variables = pm.Normal('mu_variables',mu=0, sigma=1, dims = 'group_values')
var_weights = pm.Normal("var_weights", mu = mu_variables, sigma = sigma_variables, dims = ('var_names', 'group_values'))
group_idx = pm.MutableData("group_idx", group_index, dims='obs')
print(data.shape.eval()) #[5 3]
print(var_weights.shape.eval()) #[3 2]
mu = pt.dot(data,var_weights)
print(mu.shape.eval()) #[5 2]
mu = mu[pt.arange(mu.shape.eval()[0]),group_idx] # needs to be dynamic
print(mu.shape.eval()) #[5]
mu_d = pm.Deterministic('mu_d',mu)
sigma = pm.HalfNormal('sigma', sigma=1)
pm.Normal(LKHD_LBL, mu=mu_d, sigma=sigma, observed=y)
with model:
trace = pm.sample(
draws=100,
tune=100,
chains=2,
target_accept=.99,
random_seed=123,
)
k = 2
new_data = np.random.randn(k, variables)
group_index = np.random.choice(groups, k)
with model:
pm.set_data(
coords={OBS_LBL: np.arange(k)},
new_data={
"data": new_data,
"group_idx":group_index
},
)
ppc_samples = pm.sample_posterior_predictive(
trace,
var_names=["mu_d"],
random_seed=123,
)