Metropolis-specific TypeError: The broadcastable pattern of the input is incorrect for this op

Question

I am trying to build a multilevel, multidimensional Bayesian model in PyMC3. For this question, I'll use a smaller toy model with the following graph structure:

where G represents genes, K cell types, and C_k cells of cell type k. Overall the model represents gene transcripts being sampled from a collection of cells of different cell types, where there is some binomial distribution parameterized by cell type mean expression level, mu_gk, and a cell-specific capture efficiency, p_kc.

When I sample this toy model with NUTS, it does fine and recovers reasonable posterior distributions:

import numpy as np
import pymc3 as pm
import theano.tensor as tt


# Generative model for data simulation
def sample_data(G=1, K=2, C_k=100):
    mu_gk = np.random.randint(1, 1000, size=(G, K))
    p_kc = np.random.beta(5, 95, (K, C_k))
    N_gkc = np.random.binomial(mu_gk[:, :, np.newaxis], p_kc[np.newaxis, :, :])

    return N_gkc


G = 10    # genes
K = 5     # cell types
C_k = 20  # cells per type

data = sample_data(G, K, C_k)

with pm.Model() as capture_efficiency:

    # Genes expression levels per cell type
    mu_gk = pm.Uniform('mu_gk', 1, 1000, shape=(G, K, 1))

    # Cell capture efficiencies
    p_kc = pm.Beta('p_kc', shape=(1, K, C_k), alpha=5, beta=95)

    # Captured transcripts
    N_gkc = pm.Binomial('N_gkc', shape=(G, K, C_k),
                        n=tt.tensordot(mu_gk, np.ones((C_k, 1)), [[2], [1]]),
                        p=tt.tensordot(np.ones((G, 1)), p_kc, [[1], [0]]),
                        observed=data)

    trace = pm.sample(5000, tune=10000, target_accept=0.99)

However, when I attempt to sample with Metropolis, e.g.,

    trace = pm.sample(5000, tune=10000, step=pm.Metropolis())

I receive the following stack trace and error message:

Traceback (most recent call last):
  File "/Applications/PyCharm.app/Contents/helpers/pydev/pydev_run_in_console.py", line 52, in run_file
    pydev_imports.execfile(file, globals, locals)  # execute the script
  File "/Applications/PyCharm.app/Contents/helpers/pydev/_pydev_imps/_pydev_execfile.py", line 18, in execfile
    exec(compile(contents+"\n", file, 'exec'), glob, loc)
  File "/Users/mfansler/Projects/pymc3/intro/capture-efficiency-celltypes.py", line 46, in <module>
    trace = pm.sample(5000, tune=10000,  step=pm.Metropolis(),
  File "/Users/mfansler/anaconda/envs/pymc3/lib/python3.6/site-packages/pymc3/step_methods/arraystep.py", line 65, in __new__
    step.__init__([var], *args, **kwargs)
  File "/Users/mfansler/anaconda/envs/pymc3/lib/python3.6/site-packages/pymc3/step_methods/metropolis.py", line 136, in __init__
    self.delta_logp = delta_logp(model.logpt, vars, shared)
  File "/Users/mfansler/anaconda/envs/pymc3/lib/python3.6/site-packages/pymc3/step_methods/metropolis.py", line 624, in delta_logp
    [logp0], inarray0 = pm.join_nonshared_inputs([logp], vars, shared)
  File "/Users/mfansler/anaconda/envs/pymc3/lib/python3.6/site-packages/pymc3/theanof.py", line 245, in join_nonshared_inputs
    xs_special = [theano.clone(x, replace, strict=False) for x in xs]
  File "/Users/mfansler/anaconda/envs/pymc3/lib/python3.6/site-packages/pymc3/theanof.py", line 245, in <listcomp>
    xs_special = [theano.clone(x, replace, strict=False) for x in xs]
  File "/Users/mfansler/anaconda/envs/pymc3/lib/python3.6/site-packages/theano/scan_module/scan_utils.py", line 247, in clone
    share_inputs)
  File "/Users/mfansler/anaconda/envs/pymc3/lib/python3.6/site-packages/theano/compile/pfunc.py", line 232, in rebuild_collect_shared
    cloned_v = clone_v_get_shared_updates(outputs, copy_inputs_over)
  File "/Users/mfansler/anaconda/envs/pymc3/lib/python3.6/site-packages/theano/compile/pfunc.py", line 93, in clone_v_get_shared_updates
    clone_v_get_shared_updates(i, copy_inputs_over)
  File "/Users/mfansler/anaconda/envs/pymc3/lib/python3.6/site-packages/theano/compile/pfunc.py", line 93, in clone_v_get_shared_updates
    clone_v_get_shared_updates(i, copy_inputs_over)
  File "/Users/mfansler/anaconda/envs/pymc3/lib/python3.6/site-packages/theano/compile/pfunc.py", line 93, in clone_v_get_shared_updates
    clone_v_get_shared_updates(i, copy_inputs_over)
  [Previous line repeated 9 more times]
  File "/Users/mfansler/anaconda/envs/pymc3/lib/python3.6/site-packages/theano/compile/pfunc.py", line 96, in clone_v_get_shared_updates
    [clone_d[i] for i in owner.inputs], strict=rebuild_strict)
  File "/Users/mfansler/anaconda/envs/pymc3/lib/python3.6/site-packages/theano/gof/graph.py", line 246, in clone_with_new_inputs
    new_node = self.op.make_node(*new_inputs)
  File "/Users/mfansler/anaconda/envs/pymc3/lib/python3.6/site-packages/theano/tensor/elemwise.py", line 230, in make_node
    % (self.input_broadcastable, ib)))
TypeError: The broadcastable pattern of the input is incorrect for this op. Expected (False, False, True), got (False, False, False).

I did find a GitHub issue filed for something along these lines, but I'm unclear on how the "workaround" someone proposed for their specific model would translate in my case.

I suspect the most relevant aspect of this model to the error encountered is the manual broadcasting of the parameters when instantiating the Binomial random variable:

n=tt.tensordot(mu_gk, np.ones((C_k, 1)), [[2], [1]]),
p=tt.tensordot(np.ones((G, 1)), p_kc, [[1], [0]])

which "extrudes" the 2D tensors into 3D ones matching the desired output shape.

How should this model be implemented in order to avoid the error when running Metropolis?

Answer 1

Making all parameter sampling steps flat appears to be sufficient for this toy model:

with pm.Model() as capture_efficiency:

    # Genes expression levels per cell type
    mu_gk_flat = pm.Uniform('mu_gk', 1, 1000, shape=G*K)
    mu_gk = mu_gk_flat.reshape((G, K, 1))

    # Cell capture efficiencies
    p_kc_flat = pm.Beta('p_kc', shape=K*C_k, alpha=5, beta=95)
    p_kc = p_kc_flat.reshape((1, K, C_k))

    # Captured transcripts
    N_gkc = pm.Binomial('N_gkc', shape=(G, K, C_k),
                        n=tt.tensordot(mu_gk, np.ones((C_k, 1)), [[2], [1]]),
                        p=tt.tensordot(np.ones((G, 1)), p_kc, [[1], [0]]),
                        observed=data)

    trace = pm.sample(5000, tune=10000, step=pm.Metropolis())

The N_gkc seems to be fine in tensor form, perhaps because it merely involves a likelihood calculation and not an actual sampling step.

From this toy example alone, it is not clear if additional hierarchical levels would also need to be flattened.