I’m trying to implement a series of models for binary adjacency matrices. Some of these models have multi-modal posteriors, so I decided to use PyMC’s SMC approach.
For several models this works like a charm, but the following Stochastic Block Model I have not been able to get to work. It can be sampled from using NUTS, and
pm.sample_prior_predictive() works fine as well, but if I use SMC I get the error
TypeError: index must be integers or a boolean mask.
aesara.printing.Print() commands, it seems as if the indexing I use to construct
edge_prob is where the issue resides, as
zp is printed fine, but
ep_print is never reached.
Is there a fundamental difference in how SMC uses the model definition compared to NUTS?
Note: I switched to the recently released PyMC 4.0 version, but it didn’t change the error.
import numpy as np import aesara.tensor as at import aesara.printing as ap import pymc as pm K = 3 # number of latent blocks n = 20 # number of nodes triu_indices = np.triu_indices(n, k=1) adj_vec = np.random.binomial(n=1, p=0.3, size=int(n*(n-1)/2)) a = b = 1 model = pm.Model() with model: pi = pm.Dirichlet('pi', a=np.ones(K), shape=K) z = pm.Categorical('z', p=pi, shape=n) m = int(K * (K + 1) / 2) # sample just the lower/upper triangle of K x K matrix \rho, # then transform the triangular matrix into the symmetric matrix \rho rho_vec = pm.Beta('rho_vec', alpha=a, beta=b, shape=m) tmp = pm.math.expand_packed_triangular(K, rho_vec) tmp += at.transpose(tmp) rho = pm.Deterministic('rho', tmp - 0.5 * at.diag(at.diagonal(tmp))) z_print = ap.Print('zp')(z) edge_prob = pm.Deterministic('p', rho[(z_print.reshape((n, 1)), z_print.reshape((1, n)))][triu_indices]) ep_print = ap.Print('epp')(edge_prob) A = pm.Bernoulli('A', p=ep_print, observed=adj_vec) # pp = pm.sample_prior_predictive(samples=5) # <-- works # trace_nuts = pm.sample(cores=1) # <-- works trace_smc = pm.sample_smc(cores=1) # <-- 'TypeError: index must be integers or a boolean mask'
Thanks for your help!