Tensorflow probability help: **Irrelevant to PyMC3**

Questions
1

Hi All,
Sorry for posting here, but frustrated for a long time, and will appreciate for any help.

Consider I have observations (Y,Z) \in \mathbb{R}^2 from the following fake data:

X \in [0,1], \\Y = w_1X+b_1, \\ Z = w_2X+b_2 
X_np = np.linspace(0,1,100)
# w1=2, b1=1.7, w2=3.2, b2= -0.2
Y_np = 2*X_np + 1.7
Z_np = 3.2*X_np - 0.2

I want to infer (w_1,w_2, b_1, b_2). Using PyMC3, the task was super easy since we have an observed argument in building a model.

But in Tensorflow probability I have difficulties due batch/event/sample shapes issues I dealing for weeks now.

My issues:

  • Is the code implementation corresponds to the mathematical description I gave below, in terms of shape sizes.

  • How I can calculate the target_log_prob_fn function? Is it suppose to give me a scalar or two values (since I have two observations).

Here is my implementation:

w_1 \sim \mathcal{N}(0,1) \\ w_2 \sim \mathcal{N}(0,1) \\ b_1 \sim \mathcal{N}(0,1) \\ b_2 \sim \mathcal{N}(0,1) \\ Y \sim \mathcal{N}(\mu = b_1+ w_1X, \sigma=1)\\ Z \sim \mathcal{N}(\mu = b_2+ w_2X, \sigma=1) 
Root = tfd.JointDistributionCoroutine.Root  # Convenient alias.
def model():
    b2 = yield Root(tfd.Normal(loc=tf.cast(0, dtype), scale=1.))    
    w2 = yield Root(tfd.Normal(loc=tf.cast(0, dtype), scale=1.))
    b1 = yield Root(tfd.Normal(loc=tf.cast(0, dtype), scale=1.))
    w1 = yield Root(tfd.Normal(loc=tf.cast(0, dtype), scale=1.))
    yhat = b1[...,tf.newaxis]+w1[...,tf.newaxis]*X_np
    zhat = b2[...,tf.newaxis]+w2[...,tf.newaxis]*X_np
    
    likelihood = yield tfd.Independent(tfd.Normal(loc=[yhat,zhat], scale=1.),
                                       reinterpreted_batch_ndims=2)

mdl_ols_coroutine = tfd.JointDistributionCoroutine(model)
target_log_prob_fn = lambda *x: mdl_ols_coroutine.log_prob(x + (Y_np,Z_np ))
2

I think in your case it is much easier to use tfd.JointDistributionCoroutineAutoBatched that batch a model for you, it goes something like:

@tfd.JointDistributionCoroutineAutoBatched
def mdl_ols_coroutine():
    b2 = yield tfd.Normal(loc=tf.cast(0, dtype), scale=1.)   
    w2 = yield tfd.Normal(loc=tf.cast(0, dtype), scale=1.)
    b1 = yield tfd.Normal(loc=tf.cast(0, dtype), scale=1.)
    w1 = yield tfd.Normal(loc=tf.cast(0, dtype), scale=1.)
    yhat = b1 + w1 * X_np
    zhat = b2 + w2 * X_np
    likelihood1 = yield tfd.Normal(loc=yhat, scale=1.)
    likelihood2 = yield tfd.Normal(loc=yhat, scale=1.)

val = mdl_ols_coroutine.sample(5)
mdl_ols_coroutine.log_prob_parts(val)  # <= validate shape is correct
target_log_prob_fn = lambda *x: mdl_ols_coroutine.log_prob(*x, Y_np, Z_np)
target_log_prob_fn(*val[:-2])

Also, FYI tensorflow_probability also have a user forum: https://groups.google.com/a/tensorflow.org/forum/#!forum/tfprobability

3 Likes
3

Thank you very much.

I have asked this question on tensorflow github issue section several months ago. Yesterday I figured it out there is a Google group as well and they helping me out.

Again, thanks mate.

1 Like