Theano custom Op for PyMC3

Question

I tried to implement a custom Op to which I get the "Bad input argument to theano function" error. Here is the code. The problem as I understand is: how to convert the PyMC3 variables into theano understandable type?

import numpy as np
import theano
import theano.tensor as t
from theano import config
config.compute_test_value = 'off'

#true_Data = [1,2]
#values=[]

class trial_Op(theano.Op):
    __props__ = ()
    itypes = [t.dmatrix, t.dmatrix, t.dmatrix]
    otypes = [t.dmatrix]

    def perform(self,node,inputs,output_storage):
        x0 = inputs[0]
        x1 = inputs[1]
        x2 = inputs[2]
        z = output_storage[0]
        z[0] = np.add(x0,x1)
        z[0] = np.add(z[0],x2)

    def grad(self,inputs,output_grads):
        return output_grads[0]
Trial_Op = trial_Op()

x1 = t.dmatrix()
x2 = t.dmatrix()
x3 = t.dmatrix()
f = theano.function( [x1,x2,x3], trial_Op()(x1,x2,x3) )

# the Op works for the 
#inp1 = np.random.rand(3,1)  # a 2d matrix
#inp2 = np.random.rand(3,1)  # a 2d matrix
#inp3 = np.array([[-40]])                # a constant
#print("Op application gives = ", f(inp1,inp2,inp3))


import pymc3 as pm
true_Data = [[1]]

with pm.Model() as model:
    x1 = pm.Normal('x1', mu = 0, sd = 0.1)
    x2 = pm.Normal('x2', mu = 3, sd = 0.5)
    x3 = np.asarray([[4]], dtype='float64')
#    x1 = x1.reshape(1,1)
#    x2 = x2.reshape(1,1)
    sum_of_x1_x2_x3 = f(x1,x2,x3)
    z = pm.Normal('z', sum_of_x1_x2_x3, observed = true_Data)
    start = {'x1':[[0.1]], 'x2':[[0.1]]}
    step = pm.Metropolis()
    trace = pm.sample(100, step, start)

pm.traceplot(trace)

sejabs · Answer 1 · 2017-04-28T11:59:37.797

I think I can answer your question.

Firstly, you should not use f = theano.function( [x1,x2,x3], trial_Op()(x1,x2,x3) ). Once it defined, f takes numeric values as arguments. However in the pymc3 model, x1 and x2 defined as Normal are not numeric but symbolic. So it will throw the error you just faced. If you are familiar with @as_op method documented in the tutorial, the solution is simple: change sum_of_x1_x2_x3 = f(x1,x2,x3) to sum_of_x1_x2_x3 = Trial_Op(x1,x2,x3) ;

secondly, in your code, it seems dmatrix type is not necessary. so I modify the code following:

N = 20 #data array length

class Trial_Op(theano.Op):
    __props__ = ()
    itypes = [t.dscalar, t.dscalar]
    otypes = [t.dvector] #if the data has multiple values i.e. data array

    def perform(self,node,inputs,output_storage):
        x0 = inputs[0]
        x1 = inputs[1]

        f = np.add(x0,x1)
        out = np.empty(N)
        out[:] = f
        z = output_storage[0]
        z[0] = out


trial_Op = Trial_Op()


import pymc3 as pm
true_Data = np.random.normal(2,1,N)

with pm.Model() as model:
    x1 = pm.Normal('x1', mu = 0, sd = 0.1)
    x2 = pm.Normal('x2', mu = 3, sd = 0.5)    

    mu = trial_Op(x1, x2)
    z = pm.Normal('z', mu = mu, sd = 1., observed = true_Data)

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

pm.traceplot(trace)

Note in the custom function, the otypes is dvector to satisfy the mu in the z = pm.Normal('z', mu = mu, sd = 1., observed = true_Data). And the sample number expands to 10000. The result image shows:

result

However, I don't know how to define the grad() method in the custom theano function. Maybe later someone or me can solve it to enable NUTS sampling method in the model :).

Theano custom Op for PyMC3

1 Answers1