Variables converted to shared variables by assigning step methods do not have broadcastable dimensions identified

[ctm22396]

I was running a model that required broadcasting a variable along a length 1 dimension. For example:

import numpy as np
import pymc3 as pm

y = np.random.choice([0, 1], size=(100, 40))

with pm.Model() as model:
    a = pm.Lognormal('a', shape=(1, 40))
    b = pm.Normal('b', shape=(1, 40))
    
    theta = pm.Normal('theta', shape=(100, 1))
    logit_p = pm.math.dot(theta, a) + b

    yhat = pm.Bernoulli('yhat', logit_p=logit_p, observed=y)

    trace = pm.sample(50, tune=0, chains=1,
              step = pm.Metropolis(vars=[theta]))

Assigning a separate step method to a, b and theta (or rather setting a step method to theta and letting a, b be auto-assigned) would yield the following error:

TypeError: Type TensorType(float64, matrix) (of Variable theta_shared__) \
into Type TensorType(float64, col) ...[truncated]

The issue comes from when a, b, and theta are transformed into a shared variable when either initializing the step method (in the call to make_shared_replacements made by Metropolis) or initializing the ValueGradFunction object (called by NUTS), the broadcastable argument is not provided. By default theano shared variables are not broadcastable. So without setting the broadcastable argument, the dot on theta and a fails.

This can be solved by setting the theano.shared broadcastable argument equal to var.broadcastable. That way, the shared replacement for the has the same broadcasting rules as the pymc variable it is standing in for. If the variable is scalar, the broadcastable argument cannot be set, so we must split the cases as follows:

Example from theanof.py make_shared_replacements: Instead of

return {var: theano.shared(var.tag.test_value, var.name + '_shared') for var in othervars}

I propose

out = {}
for var in othervars:
    broadcast = var.broadcastable
    if len(broadcast):
        out[var] = theano.shared(var.tag.test_value, var.name + '_shared',
                                 broadcastable=broadcast) # <- this is the solution
    else:
        out[var] = theano.shared(var.tag.test_value, var.name + '_shared')
return out

and in model.py __init__ of ValueGradFunction: Instead of

for var in extra_vars:
    shared = theano.shared(var.tag.test_value, var.name + '_shared__')
    self._extra_vars_shared[var.name] = shared

I propose

for var in extra_vars:
    broadcast = var.broadcastable
    if len(broadcast):
        shared = theano.shared(var.tag.test_value,
                               var.name + '_shared__',
                               broadcastable=broadcast)
    else:
        shared = theano.shared(var.tag.test_value,
                               var.name + '_shared__')

    self._extra_vars_shared[var.name] = shared

And likewise wherever else this problem arises.

Here is additional information about my environment:

• PyMC3 Version: 3.6
• Theano Version: 1.0.2
• Python Version: 3.7.1
• Operating system: linux
• How did you install PyMC3: conda

[michaelosthege]

closed by #4492