Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Support TensorFlow Probability #20

Closed
brandonwillard opened this issue Jun 1, 2019 · 2 comments
Closed

Support TensorFlow Probability #20

brandonwillard opened this issue Jun 1, 2019 · 2 comments
Labels
enhancement New feature or request help wanted Extra attention is needed important Features and issues that need to be addressed ASAP TensorFlow This issue involves the TensorFlow backend

Comments

@brandonwillard
Copy link
Contributor

We need meta object and S-expression support for TensorFlow Probability distributions/classes.
@brandonwillard brandonwillard added important Features and issues that need to be addressed ASAP enhancement New feature or request help wanted Extra attention is needed and removed important Features and issues that need to be addressed ASAP labels Jun 1, 2019
@brandonwillard
Copy link
Contributor Author

brandonwillard commented Jul 2, 2019 •
edited
Loading

It looks like we might already have almost everything we need. From the example in our unit tests:

import numpy as np

import tensorflow as tf

from tensorflow_probability import distributions as tfd

from symbolic_pymc.tensorflow.meta import mt

from symbolic_pymc.tensorflow.printing import tf_dprint

N = 100
sigma_tf = tfd.Gamma(np.asarray(1.), np.asarray(1.), name='sigma').sample()
epsilon_tf = tfd.Normal(np.zeros((N, 1)), sigma_tf, name='epsilon').sample()
beta_tf = tfd.Normal(np.zeros((2, 1)), 1, name='beta').sample()

X = np.vstack([np.random.randn(N), np.ones(N)]).T
X_tf = tf.convert_to_tensor(X)

Y_tf = tf.linalg.matmul(X_tf, beta_tf) + epsilon_tf

we get the following graph, which appears to represent the entire sample space (e.g. distributions, their dependencies, operations on samples, etc.):

>>> tf_dprint(Y_tf)
Tensor(Add):0,	shape=[100, 1]	"add:0"
|  Op(Add)	"add"
|  |  Tensor(MatMul):0,	shape=[100, 1]	"MatMul:0"
|  |  |  Op(MatMul)	"MatMul"
|  |  |  |  Tensor(Const):0,	shape=[100, 2]	"Const:0"
|  |  |  |  Tensor(Reshape):0,	shape=[2, 1]	"beta_1/sample/Reshape:0"
|  |  |  |  |  Op(Reshape)	"beta_1/sample/Reshape"
|  |  |  |  |  |  Tensor(Add):0,	shape=[1, 2, 1]	"beta_1/sample/add:0"
|  |  |  |  |  |  |  Op(Add)	"beta_1/sample/add"
|  |  |  |  |  |  |  |  Tensor(Mul):0,	shape=[1, 2, 1]	"beta_1/sample/mul:0"
|  |  |  |  |  |  |  |  |  Op(Mul)	"beta_1/sample/mul"
|  |  |  |  |  |  |  |  |  |  Tensor(Add):0,	shape=[1, 2, 1]	"beta_1/sample/random_normal:0"
|  |  |  |  |  |  |  |  |  |  |  Op(Add)	"beta_1/sample/random_normal"
|  |  |  |  |  |  |  |  |  |  |  |  Tensor(Mul):0,	shape=[1, 2, 1]	"beta_1/sample/random_normal/mul:0"
|  |  |  |  |  |  |  |  |  |  |  |  |  Op(Mul)	"beta_1/sample/random_normal/mul"
|  |  |  |  |  |  |  |  |  |  |  |  |  |  Tensor(RandomStandardNormal):0,	shape=[1, 2, 1]	"beta_1/sample/random_normal/RandomStandardNormal:0"
|  |  |  |  |  |  |  |  |  |  |  |  |  |  |  Op(RandomStandardNormal)	"beta_1/sample/random_normal/RandomStandardNormal"
|  |  |  |  |  |  |  |  |  |  |  |  |  |  |  |  Tensor(ConcatV2):0,	shape=[3]	"beta_1/sample/concat:0"
|  |  |  |  |  |  |  |  |  |  |  |  |  |  |  |  |  Op(ConcatV2)	"beta_1/sample/concat"
|  |  |  |  |  |  |  |  |  |  |  |  |  |  |  |  |  |  Tensor(Const):0,	shape=[1]	"beta_1/sample/concat/values_0:0"
|  |  |  |  |  |  |  |  |  |  |  |  |  |  |  |  |  |  Tensor(Identity):0,	shape=[2]	"beta_1/sample/beta/batch_shape_tensor/batch_shape:0"
|  |  |  |  |  |  |  |  |  |  |  |  |  |  |  |  |  |  |  Op(Identity)	"beta_1/sample/beta/batch_shape_tensor/batch_shape"
|  |  |  |  |  |  |  |  |  |  |  |  |  |  |  |  |  |  |  |  Tensor(Const):0,	shape=[2]	"beta_1/sample/beta/batch_shape_tensor/Const:0"
|  |  |  |  |  |  |  |  |  |  |  |  |  |  |  |  |  |  Tensor(Const):0,	shape=[]	"beta_1/sample/concat/axis:0"
|  |  |  |  |  |  |  |  |  |  |  |  |  |  Tensor(Const):0,	shape=[]	"beta_1/sample/random_normal/stddev:0"
|  |  |  |  |  |  |  |  |  |  |  |  Tensor(Const):0,	shape=[]	"beta_1/sample/random_normal/mean:0"
|  |  |  |  |  |  |  |  |  |  Tensor(Const):0,	shape=[]	"beta/scale:0"
|  |  |  |  |  |  |  |  Tensor(Const):0,	shape=[2, 1]	"beta/loc:0"
|  |  |  |  |  |  Tensor(ConcatV2):0,	shape=[2]	"beta_1/sample/concat_1:0"
|  |  |  |  |  |  |  Op(ConcatV2)	"beta_1/sample/concat_1"
|  |  |  |  |  |  |  |  Tensor(Const):0,	shape=[0]	"beta_1/sample/sample_shape:0"
|  |  |  |  |  |  |  |  Tensor(StridedSlice):0,	shape=[2]	"beta_1/sample/strided_slice:0"
|  |  |  |  |  |  |  |  |  Op(StridedSlice)	"beta_1/sample/strided_slice"
|  |  |  |  |  |  |  |  |  |  Tensor(Const):0,	shape=[3]	"beta_1/sample/Shape:0"
|  |  |  |  |  |  |  |  |  |  Tensor(Const):0,	shape=[1]	"beta_1/sample/strided_slice/stack:0"
|  |  |  |  |  |  |  |  |  |  Tensor(Const):0,	shape=[1]	"beta_1/sample/strided_slice/stack_1:0"
|  |  |  |  |  |  |  |  |  |  Tensor(Const):0,	shape=[1]	"beta_1/sample/strided_slice/stack_2:0"
|  |  |  |  |  |  |  |  Tensor(Const):0,	shape=[]	"beta_1/sample/concat_1/axis:0"
|  |  Tensor(Reshape):0,	shape=[100, 1]	"epsilon_1/sample/Reshape:0"
|  |  |  Op(Reshape)	"epsilon_1/sample/Reshape"
|  |  |  |  Tensor(Add):0,	shape=[1, 100, 1]	"epsilon_1/sample/add:0"
|  |  |  |  |  Op(Add)	"epsilon_1/sample/add"
|  |  |  |  |  |  Tensor(Mul):0,	shape=[1, 100, 1]	"epsilon_1/sample/mul:0"
|  |  |  |  |  |  |  Op(Mul)	"epsilon_1/sample/mul"
|  |  |  |  |  |  |  |  Tensor(Add):0,	shape=[1, 100, 1]	"epsilon_1/sample/random_normal:0"
|  |  |  |  |  |  |  |  |  Op(Add)	"epsilon_1/sample/random_normal"
|  |  |  |  |  |  |  |  |  |  Tensor(Mul):0,	shape=[1, 100, 1]	"epsilon_1/sample/random_normal/mul:0"
|  |  |  |  |  |  |  |  |  |  |  Op(Mul)	"epsilon_1/sample/random_normal/mul"
|  |  |  |  |  |  |  |  |  |  |  |  Tensor(RandomStandardNormal):0,	shape=[1, 100, 1]	"epsilon_1/sample/random_normal/RandomStandardNormal:0"
|  |  |  |  |  |  |  |  |  |  |  |  |  Op(RandomStandardNormal)	"epsilon_1/sample/random_normal/RandomStandardNormal"
|  |  |  |  |  |  |  |  |  |  |  |  |  |  Tensor(ConcatV2):0,	shape=[3]	"epsilon_1/sample/concat:0"
|  |  |  |  |  |  |  |  |  |  |  |  |  |  |  Op(ConcatV2)	"epsilon_1/sample/concat"
|  |  |  |  |  |  |  |  |  |  |  |  |  |  |  |  Tensor(Const):0,	shape=[1]	"epsilon_1/sample/concat/values_0:0"
|  |  |  |  |  |  |  |  |  |  |  |  |  |  |  |  Tensor(Identity):0,	shape=[2]	"epsilon_1/sample/epsilon/batch_shape_tensor/batch_shape:0"
|  |  |  |  |  |  |  |  |  |  |  |  |  |  |  |  |  Op(Identity)	"epsilon_1/sample/epsilon/batch_shape_tensor/batch_shape"
|  |  |  |  |  |  |  |  |  |  |  |  |  |  |  |  |  |  Tensor(Const):0,	shape=[2]	"epsilon_1/sample/epsilon/batch_shape_tensor/Const:0"
|  |  |  |  |  |  |  |  |  |  |  |  |  |  |  |  Tensor(Const):0,	shape=[]	"epsilon_1/sample/concat/axis:0"
|  |  |  |  |  |  |  |  |  |  |  |  Tensor(Const):0,	shape=[]	"epsilon_1/sample/random_normal/stddev:0"
|  |  |  |  |  |  |  |  |  |  Tensor(Const):0,	shape=[]	"epsilon_1/sample/random_normal/mean:0"
|  |  |  |  |  |  |  |  Tensor(Reshape):0,	shape=[]	"sigma_2/sample/Reshape:0"
|  |  |  |  |  |  |  |  |  Op(Reshape)	"sigma_2/sample/Reshape"
|  |  |  |  |  |  |  |  |  |  Tensor(Maximum):0,	shape=[1]	"sigma_2/sample/random_gamma/Maximum:0"
|  |  |  |  |  |  |  |  |  |  |  Op(Maximum)	"sigma_2/sample/random_gamma/Maximum"
|  |  |  |  |  |  |  |  |  |  |  |  Tensor(Const):0,	shape=[]	"sigma_2/sample/random_gamma/Maximum/x:0"
|  |  |  |  |  |  |  |  |  |  |  |  Tensor(RealDiv):0,	shape=[1]	"sigma_2/sample/random_gamma/truediv:0"
|  |  |  |  |  |  |  |  |  |  |  |  |  Op(RealDiv)	"sigma_2/sample/random_gamma/truediv"
|  |  |  |  |  |  |  |  |  |  |  |  |  |  Tensor(RandomGamma):0,	shape=[1]	"sigma_2/sample/random_gamma/RandomGamma:0"
|  |  |  |  |  |  |  |  |  |  |  |  |  |  |  Op(RandomGamma)	"sigma_2/sample/random_gamma/RandomGamma"
|  |  |  |  |  |  |  |  |  |  |  |  |  |  |  |  Tensor(Const):0,	shape=[1]	"sigma_2/sample/random_gamma/shape:0"
|  |  |  |  |  |  |  |  |  |  |  |  |  |  |  |  Tensor(Add):0,	shape=[]	"sigma_2/sample/random_gamma/add:0"
|  |  |  |  |  |  |  |  |  |  |  |  |  |  |  |  |  Op(Add)	"sigma_2/sample/random_gamma/add"
|  |  |  |  |  |  |  |  |  |  |  |  |  |  |  |  |  |  Tensor(Identity):0,	shape=[]	"sigma/Identity:0"
|  |  |  |  |  |  |  |  |  |  |  |  |  |  |  |  |  |  |  Op(Identity)	"sigma/Identity"
|  |  |  |  |  |  |  |  |  |  |  |  |  |  |  |  |  |  |  |  Tensor(Const):0,	shape=[]	"sigma/concentration:0"
|  |  |  |  |  |  |  |  |  |  |  |  |  |  |  |  |  |  Tensor(Const):0,	shape=[]	"sigma_2/sample/random_gamma/zeros_like:0"
|  |  |  |  |  |  |  |  |  |  |  |  |  |  Tensor(Identity):0,	shape=[]	"sigma/Identity_1:0"
|  |  |  |  |  |  |  |  |  |  |  |  |  |  |  Op(Identity)	"sigma/Identity_1"
|  |  |  |  |  |  |  |  |  |  |  |  |  |  |  |  Tensor(Const):0,	shape=[]	"sigma/rate:0"
|  |  |  |  |  |  |  |  |  |  Tensor(ConcatV2):0,	shape=[0]	"sigma_2/sample/concat:0"
|  |  |  |  |  |  |  |  |  |  |  Op(ConcatV2)	"sigma_2/sample/concat"
|  |  |  |  |  |  |  |  |  |  |  |  Tensor(Const):0,	shape=[0]	"sigma_2/sample/sample_shape:0"
|  |  |  |  |  |  |  |  |  |  |  |  Tensor(StridedSlice):0,	shape=[0]	"sigma_2/sample/strided_slice:0"
|  |  |  |  |  |  |  |  |  |  |  |  |  Op(StridedSlice)	"sigma_2/sample/strided_slice"
|  |  |  |  |  |  |  |  |  |  |  |  |  |  Tensor(Const):0,	shape=[1]	"sigma_2/sample/Shape:0"
|  |  |  |  |  |  |  |  |  |  |  |  |  |  Tensor(Const):0,	shape=[1]	"sigma_2/sample/strided_slice/stack:0"
|  |  |  |  |  |  |  |  |  |  |  |  |  |  Tensor(Const):0,	shape=[1]	"sigma_2/sample/strided_slice/stack_1:0"
|  |  |  |  |  |  |  |  |  |  |  |  |  |  Tensor(Const):0,	shape=[1]	"sigma_2/sample/strided_slice/stack_2:0"
|  |  |  |  |  |  |  |  |  |  |  |  Tensor(Const):0,	shape=[]	"sigma_2/sample/concat/axis:0"
|  |  |  |  |  |  Tensor(Const):0,	shape=[100, 1]	"epsilon/loc:0"
|  |  |  |  Tensor(ConcatV2):0,	shape=[2]	"epsilon_1/sample/concat_1:0"
|  |  |  |  |  Op(ConcatV2)	"epsilon_1/sample/concat_1"
|  |  |  |  |  |  Tensor(Const):0,	shape=[0]	"epsilon_1/sample/sample_shape:0"
|  |  |  |  |  |  Tensor(StridedSlice):0,	shape=[2]	"epsilon_1/sample/strided_slice:0"
|  |  |  |  |  |  |  Op(StridedSlice)	"epsilon_1/sample/strided_slice"
|  |  |  |  |  |  |  |  Tensor(Const):0,	shape=[3]	"epsilon_1/sample/Shape:0"
|  |  |  |  |  |  |  |  Tensor(Const):0,	shape=[1]	"epsilon_1/sample/strided_slice/stack:0"
|  |  |  |  |  |  |  |  Tensor(Const):0,	shape=[1]	"epsilon_1/sample/strided_slice/stack_1:0"
|  |  |  |  |  |  |  |  Tensor(Const):0,	shape=[1]	"epsilon_1/sample/strided_slice/stack_2:0"
|  |  |  |  |  |  Tensor(Const):0,	shape=[]	"epsilon_1/sample/concat_1/axis:0"

By converting to a meta graph, as we have been, those distribution tf.Operations (e.g.
Op(RandomStandardNormal), Op(RandomGamma)) are already covered:

>>> Y_mt = mt(Y_tf)
TFlowMetaTensor(
  dtype=tf.float64,
  op=TFlowMetaOp(
    op_def=TFlowMetaOpDef(Add),
    node_def=TFlowMetaNodeDef( op='Add', name='add', attr={}),
    inputs=(TFlowMetaTensor(
       dtype=tf.float64,
       op=TFlowMetaOp(
         op_def=TFlowMetaOpDef(MatMul),
         node_def=TFlowMetaNodeDef( op='MatMul', name='MatMul', attr={}),
         inputs=(TFlowMetaConstant( ),
          TFlowMetaTensor(
            dtype=tf.float64,
            op=TFlowMetaOp(
              op_def=TFlowMetaOpDef(Reshape),
              node_def=TFlowMetaNodeDef(
                op='Reshape',
                name='beta_1/sample/Reshape',
                attr={}),
              inputs=(TFlowMetaTensor(
                 dtype=tf.float64,
                 op=TFlowMetaOp(
                   op_def=TFlowMetaOpDef(Add),
                   node_def=TFlowMetaNodeDef(
                     op='Add',
                     name='beta_1/sample/add',
                     attr={}),
                   inputs=(TFlowMetaTensor(
                      dtype=tf.float64,
                      op=TFlowMetaOp(
                        op_def=TFlowMetaOpDef(Mul),
                        node_def=TFlowMetaNodeDef(
                          op='Mul',
                          name='beta_1/sample/mul',
                          attr={}),
                        inputs=(TFlowMetaTensor(
                           dtype=tf.float64,
                           op=TFlowMetaOp(
                             op_def=TFlowMetaOpDef(Add),
                             node_def=TFlowMetaNodeDef(
                               op='Add',
                               name='beta_1/sample/random_normal',
                               attr={}),
                             inputs=(TFlowMetaTensor(
                                dtype=tf.float64,
                                op=TFlowMetaOp(
                                  op_def=TFlowMetaOpDef(Mul),
                                  node_def=TFlowMetaNodeDef(
                                    op='Mul',
                                    name='beta_1/sample/random_normal/mul',
                                    attr={}),
                                  inputs=(TFlowMetaTensor(
                                     dtype=tf.float64,
                                     op=TFlowMetaOp(
                                       op_def=TFlowMetaOpDef(RandomStandardNormal),
                                       node_def=TFlowMetaNodeDef(
                                         op='RandomStandardNormal',
                                         name='beta_1/sample/random_normal/RandomStandardNormal',
                                         attr={'dtype': 'float64'}),
                                       inputs=(TFlowMetaTensor(...

In other words, we should now be able to create all the same relations we have in Theano in TensorFlow (e.g. recenter/rescale, conjugation, etc.), without the need to create more meta objects or features/support! We simply need to target those specific TFP ops and the formats of their arguments.

@brandonwillard
Copy link
Contributor Author

Currently, the graphs produced by TFP are fully amenable to symbolic-pymc manipulation. For ease of use, we should provide graph normalization goals so that it's easier to manipulate the output of TFP and convert TF graphs to TFP objects.

That work is being addressed in #19, so I'm closing this one.

@brandonwillard brandonwillard added the TensorFlow This issue involves the TensorFlow backend label Mar 12, 2020
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Assignees
No one assigned
Labels
enhancement New feature or request help wanted Extra attention is needed important Features and issues that need to be addressed ASAP TensorFlow This issue involves the TensorFlow backend
Projects
None yet
Milestone
No milestone
Development

No branches or pull requests
1 participant
@brandonwillard