Model works with NUTS/prior predictive, but not with SMC?


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!

@ricardoV94 could you please take a look?

I really want to look at this issue, just haven’t got the time yet!

No worries :slight_smile:

Might the issue be related to SMC not using discrete proposals within the sampling chains?

We don’t use discrete proposals but we round it for the case of discrete variables

So it seems SMC is failing with a simple categorical model:

with pm.Model() as m:
    mu = pm.Categorical("mu", p=[0.1, 0.3, 0.6], size=2)
    pm.Normal("like", mu=mu, sigma=0.1, observed=[1, 2])

    trace = pm.sample_smc(chains=1, return_inferencedata=False)

This PR should fix it: Cast rounded input replacement to original discrete dtype in sample_smc by ricardoV94 · Pull Request #5887 · pymc-devs/pymc · GitHub

1 Like

Awesome, thank you so much!

Unfortunately, although that PR fixes the TypeError, this is now replaced by errors like these: IndexError: index 3 is out of bounds for axis 0 with size 3.

When using the z_print = ap.Print('zp')(z) approach, it shows that among the draws, occasionally z is outside of the options dictated by the Categorical - whereas my options are [0, …, K-1], some values for z[i] are -1 or K. Should perhaps the rounding be replaced by a ceil() or floor(), depending on what float input is sampled?

That’s a more structural problem with the logp evaluation:

import pymc as pm
import aesara.tensor as at
with pm.Model() as m:
    x = pm.Categorical("x", p=[0.1, 0.3, 0.6])
    y = pm.Normal("y", mu=at.constant([0, 1, 2])[x])
m.compile_logp()({"x": 4, "y": 0})  # IndexError

I think by default the Categorical in pm.sample is given a step sampler that never proposes invalid values, but otherwise we don’t have guards for those cases. This idea would help with that in the future: Add some transforms to discrete variables · Discussion #5888 · pymc-devs/pymc · GitHub

For the time being you could clip x explicitly like in:

import pymc as pm
import aesara.tensor as at
with pm.Model() as m:
    x = pm.Categorical("x", p=[0.1, 0.3, 0.6])
    y = pm.Normal("y", mu=at.constant([0, 1, 2])[at.clip(x, 0, 2)])
m.compile_logp()({"x": 4, "y": 0})  # -inf, no error!

Thank you again, this fixes it for me! :slight_smile: