I have the following pymc3 (3.8) hierarchical model, where
- y (dependent variable) is a 1313x1
- x (level 1 predictors) is a 1313x2
- z (level 2 predictors) is a 1313x3 matrix
with pm.Model() as model:
gamma_0_0 = pm.Normal("gamma_0_0", mu=0, sigma=10)
gamma_k_0 = pm.Normal("gamma_k_0", mu=0, sigma=10, shape=x.shape[1])
gamma_0_j = pm.Normal("gamma_0_j", mu=0, sigma=10, shape=z.shape[1])
gamma_k_j = pm.Normal("gamma_k_j", mu=0, sigma=10, shape=(z.shape[1], x.shape[1]))
# Model error
u_0_j = pm.Normal("u_0_j", mu=0, sigma=10)
u_k_j = pm.Normal("u_k_j", mu=0, sigma=10, shape=x.shape[1])
r_ij = pm.Gamma("r_ij", alpha=9, beta=4)
beta_0_j = pm.math.dot(z, gamma_0_j) + gamma_0_0 + u_0_j
beta_k_j = pm.math.dot(z, gamma_k_j) + gamma_k_0 + u_k_j
y_hat = pm.math.dot(x, beta_k_j.T) + beta_0_j
# Likelihood
y_like = pm.Normal("y_like", mu=y_hat, sigma=r_ij, observed=y)
with model:
step = pm.NUTS()
trace = pm.sample(draws=500, chains=2, tune=200, discard_tuned_samples=True, random_seed=SEED, step=step)
When I draw only 200 samples (tune=100) from the posterior, the model works fine. However, when I increase the samples to 500 (tune=200), the following error comes up:
numpy.core._exceptions.MemoryError: Unable to allocate 12.8 GiB for an array with shape (2, 500, 1313, 1313) and data type float64
Absent the issue that the model does no longer runs, I don't understand why the array shape in the error report is (2, 500, 1313, 1313). 2 could be the number of chains, 500 the number of samples, and 1313 the length of y. But why is it not (2, 500, 1313)?
I know that this error can probably be worked around by allocating more memory as explained here (Unable to allocate array with shape and data type) but I am seeking a more sustainable solution.
I use Windows10 (32GB RAM), Anaconda 3 and Python 3.8