Custom distribution using sympy

Questions
2

The functions sympy.printing.theanocode.[theano_code, theano_function] are sufficient for converting straightforward SymPy graphs to Theano graphs.

The primary difficulty I’ve noticed involves the re-use of existing Theano variables (e.g. the ones produced by PyMC3).

Here’s an example:

import theano
import theano.tensor as tt

import sympy as sp
import sympy.stats as st

from theano.printing import debugprint as tt_dprint
from sympy.printing.theanocode import theano_code


# theano.config.cxx = ""
# theano.config.mode = "FAST_COMPILE"
# tt.config.compute_test_value = 'ignore'


k_sp, theta_sp, y_sp = sp.symbols(["k", "theta", "y"])

X_sp = st.Gamma("X", k_sp, theta_sp)

gamma_pdf_sp = st.density(X_sp, y_sp)


# Create a Theano graph for the SymPy PDF graph
gamma_pdf_tt = theano_code(gamma_pdf_sp(y_sp))

tt_dprint(gamma_pdf_tt)

The output is

Elemwise{mul,no_inplace} [id A] ''   
 |Elemwise{pow,no_inplace} [id B] ''   
 | |theta [id C]
 | |Elemwise{mul,no_inplace} [id D] ''   
 |   |TensorConstant{-1} [id E]
 |   |k [id F]
 |Elemwise{pow,no_inplace} [id G] ''   
 | |y [id H]
 | |Elemwise{add,no_inplace} [id I] ''   
 |   |TensorConstant{-1} [id E]
 |   |k [id F]
 |Elemwise{pow,no_inplace} [id J] ''   
 | |Elemwise{gamma,no_inplace} [id K] ''   
 | | |k [id F]
 | |TensorConstant{-1} [id E]
 |Elemwise{exp,no_inplace} [id L] ''   
   |Elemwise{mul,no_inplace} [id M] ''   
     |TensorConstant{-1} [id E]
     |y [id H]
     |Elemwise{pow,no_inplace} [id N] ''   
       |theta [id C]
       |TensorConstant{-1} [id E]

Now, let’s say we had our own Theano variable that we wanted to use in place of the SymPy variable k_sp (e.g. the kind produced by pm.Normal, etc.):

def create_cache_key(s, name=None, dtype=None, broadcastable=None):
    """Create a TheanoPrinter cache key for a SymPy variable."""
    if name is None:
        name = s.name
    if dtype is None:
        dtype = 'floatX'
    if broadcastable is None:
        broadcastable = ()
    return (name, type(s), s.args, dtype, broadcastable)



# An existing Theano variable that we want to use in the output Theano graph
k_tt = tt.scalar('k_ours')

gamma_pdf_tt = theano_code(gamma_pdf_sp(y_sp), cache={create_cache_key(k_sp): k_tt})

tt_dprint(gamma_pdf_tt)

Now, the output Theano graph contains our k_ours variable:

Elemwise{mul,no_inplace} [id A] ''   
 |Elemwise{pow,no_inplace} [id B] ''   
 | |theta [id C]
 | |Elemwise{mul,no_inplace} [id D] ''   
 |   |TensorConstant{-1} [id E]
 |   |k_ours [id F]
 |Elemwise{pow,no_inplace} [id G] ''   
 | |y [id H]
 | |Elemwise{add,no_inplace} [id I] ''   
 |   |TensorConstant{-1} [id E]
 |   |k_ours [id F]
 |Elemwise{pow,no_inplace} [id J] ''   
 | |Elemwise{gamma,no_inplace} [id K] ''   
 | | |k_ours [id F]
 | |TensorConstant{-1} [id E]
 |Elemwise{exp,no_inplace} [id L] ''   
   |Elemwise{mul,no_inplace} [id M] ''   
     |TensorConstant{-1} [id E]
     |y [id H]
     |Elemwise{pow,no_inplace} [id N] ''   
       |theta [id C]
       |TensorConstant{-1} [id E]

The same cache={...} approach should work for theano_function, as well.

1 Like