Hi,
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
.
Through some 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!