Apply to speak at TensorFlow World. Deadline April 23rd. Propose talk


    log_rate_fn=(lambda x: x),

Constructs a trainable tfd.Poisson distribution.

This function creates a Poisson distribution parameterized by log rate. Using default args, this function is mathematically equivalent to:

Y = Poisson(log_rate=matmul(W, x) + b)

  W in R^[d, n]
  b in R^d


This can be used as a Poisson regression loss.

# This example fits a poisson regression loss.
import numpy as np
import tensorflow as tf
import tensorflow_probability as tfp

# Create fictitious training data.
dtype = np.float32
n = 3000    # number of samples
x_size = 4  # size of single x
def make_training_data():
  x = np.random.randn(n, x_size).astype(dtype)
  w = np.random.randn(x_size).astype(dtype)
  b = np.random.randn(1).astype(dtype)
  true_log_rate = np.tensordot(x, w, axes=[[-1], [-1]]) + b
  y = np.random.poisson(lam=np.exp(true_log_rate)).astype(dtype)
  return y, x
y, x = make_training_data()

# Build TF graph for fitting Poisson maximum likelihood estimator.
poisson = tfp.trainable_distributions.poisson(x)
loss = -tf.reduce_mean(poisson.log_prob(y))
train_op = tf.train.AdamOptimizer(learning_rate=2.**-5).minimize(loss)
mse = tf.reduce_mean(tf.squared_difference(y, poisson.mean()))
init_op = tf.global_variables_initializer()

# Run graph 1000 times.
num_steps = 1000
loss_ = np.zeros(num_steps)   # Style: `_` to indicate result.
mse_ = np.zeros(num_steps)
with tf.Session() as sess:
  for it in xrange(loss_.size):
    _, loss_[it], mse_[it] =[train_op, loss, mse])
    if it % 200 == 0 or it == loss_.size - 1:
      print("iteration:{}  loss:{}  mse:{}".format(it, loss_[it], mse_[it]))

# ==> iteration:0    loss:37.0814208984  mse:6359.41259766
#     iteration:200  loss:1.42010736465  mse:40.7654914856
#     iteration:400  loss:1.39027583599  mse:8.77660560608
#     iteration:600  loss:1.3902695179   mse:8.78443241119
#     iteration:800  loss:1.39026939869  mse:8.78443622589
#     iteration:999  loss:1.39026939869  mse:8.78444766998


  • x: Tensor with floating type. Must have statically defined rank and statically known right-most dimension.
  • layer_fn: Python callable which takes input x and int scalar d and returns a transformation of x with shape tf.concat([tf.shape(x)[:-1], [1]], axis=0). Default value: tf.layers.dense.
  • log_rate_fn: Python callable which transforms the log_rate parameter. Takes a (batch of) length-dims vectors and returns a Tensor of same shape and dtype. Default value: lambda x: x.
  • name: A name_scope name for operations created by this function. Default value: None (i.e., "poisson").


  • poisson: An instance of tfd.Poisson.