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.