TypeError in NUTS

I have a model with a mix of continuous and discrete variables. From the documentation, it looks like pymc3 can handle that by using NUTS for the continuous variables and metropolis-hastings for the discrete ones.


def Matern52(X1, X2, alpha, rho):
    L2 = np.sum(X1 ** 2, axis=1, keepdims=True) + np.sum(X2 ** 2, axis=1, keepdims=True).T - 2 * X1 @ X2.T
    L2 = np.clip(L2, 0.0, None)
    L = np.sqrt(L2 + 1e-12)
    sqrt5_r = 5**0.5 * L
    return alpha * (1 + sqrt5_r / rho + (5/3) * L2 / rho ** 2) * np.exp(-sqrt5_r / rho)

X = np.linspace(-1, 1, 500).reshape((500, 1))
alpha = 1.0
rho = 0.3
sigma = 0.2
cov = Matern52(X, X, alpha, rho)
L = np.linalg.cholesky(cov)
np.random.seed(6422395)
f = np.random.multivariate_normal(np.zeros(len(X)), cov)
n_train = 50
n_obs = 10
np.random.seed(4)
train_inds = np.random.choice(500, size=n_train, replace=True)
X_train = X[train_inds]
noise = np.random.normal(loc=0, scale=sigma, size=n_train)
y_train = f[train_inds] + noise
X_obs1 = X_train[:n_obs]
y_obs1 = y_train[:n_obs]
X_hyp1 = X_train[n_obs:]

with pm.Model() as model:
    a = pm.Uniform('a', lower=1e-4, upper=2)
    r = pm.Uniform('r', lower=1e-4, upper=2)
    s = pm.Uniform('s', lower=1e-2, upper=1)
    probs = pm.DiscreteUniform('probs', lower=0, upper=499, shape=(n_train - n_obs, 1))

    p = tt.extra_ops.to_one_hot(probs, 500)
    Xh = tt.dot(p, X)
    Xt = tt.concatenate([X_obs1, Xh])
    R2 = tt.sum(Xt ** 2, axis=1, keepdims=True) + tt.sum(Xt ** 2, axis=1, keepdims=True).T 
    R2 -= 2 * tt.dot(Xt, Xt.T)
    R2 = R2 * (R2 > 0)
    R = tt.sqrt(R2 + 1e-12)
    sqrt5_r = 5**0.5 * R
    K = a * (1 + sqrt5_r / r + (5/3) * R2 / r ** 2) * tt.exp(-sqrt5_r / r) + tt.eye(50) * (s + 1e-4)
    L = tt.slinalg.Cholesky()(K)
    y = pm.MvNormal('y', mu=np.zeros_like(y_train), chol=L, observed=y_train)
    step1 = pm.NUTS(vars=[a, r, s])
    step2 = pm.Metropolis(vars=[probs])

    trace_discrete = pm.sample(500, step=[step1, step2], tune=1000)

When I run this, I get TypeError: Cannot convert Type TensorType(int64, matrix) (of Variable probs_shared__) into Type TensorType(int64, col). You can try to manually convert probs_shared__ into a TensorType(int64, col)..

Any ideas as to what might be causing this?

There shape error is a bit odd, but reparameterize it into something like below should work:

Xshare = theano.shared(X) # cast into theano so you can index it using pymc3/theano tensor
with pm.Model() as model:
    ...
    # use categorical because Metropolis gives invalid proposal value for DiscreteUniform sometimes
    probs = pm.Categorical('probs', p=np.ones(500), shape=n_train - n_obs)
    Xh = Xshare[probs] # tt.dot(p, X)
    ...
    # sample using the auto assigned sampler
    trace_discrete = pm.sample(500, tune=1000)

However, models with both discrete variable and continuous variable are difficult to inference (NUTS), you should try to find a way to reparameterize it further.

1 Like

That seems to work.

Right before you posted that, I figured out that this error goes away if I do

with pm.Model() as model:
    ...
    # use categorical because Metropolis gives invalid proposal value for DiscreteUniform sometimes
    probs = pm.DiscreteUniform('probs', lower=0, upper=499, shape=(n_train - n_obs))
    p = tt.extra_ops.to_one_hot(probs.reshape((n_train - n_obs, 1)), 500)
    ...
    # sample using the auto assigned sampler
    trace_discrete = pm.sample(500, tune=1000)

And it will start running, but somewhere around step 200 it throws a very long and unintelligible IndexError.

Using your method, it runs, probs is always zero in the trace – it never explores.

Oh wait, I was bad at following directions and didn’t use the auto-assigned sampler like you said. Once I did that, it worked. Thank you!

No problem.
The index error comes from the Metropolis sampler - it gives invalid proposal value for DiscreteUniform (for example propose +2 when the last value is 499, which gives the next value 501 that goes out of bound).

2 Likes