def testMakeRandomVariable(self):
"""Tests that manual wrapping is the same as the built-in solution."""
custom_normal = ed.make_random_variable(tfp.distributions.Normal)
def model_builtin():
return ed.Normal(1., 0.1, name="x")
def model_wrapped():
return custom_normal(1., 0.1, name="x")
log_joint_builtin = ed.make_log_joint_fn(model_builtin)
log_joint_wrapped = ed.make_log_joint_fn(model_wrapped)
self.assertEqual(log_joint_builtin(x=7.), log_joint_wrapped(x=7.))
def testTrivialTracerPreservesLogJoint(self):
def trivial_tracer(fn, *args, **kwargs):
# A tracer that does nothing.
return ed.traceable(fn)(*args, **kwargs)
def model():
return ed.Normal(0., 1., name="x")
def transformed_model():
with ed.trace(trivial_tracer):
model()
log_joint = ed.make_log_joint_fn(model)
log_joint_transformed = ed.make_log_joint_fn(transformed_model)
self.assertEqual(log_joint(x=5.), log_joint_transformed(x=5.))
def testMakeLogJointFnUnconditional(self):
"""Test `make_log_joint_fn` on unconditional Edward program."""
def normal_with_unknown_mean():
loc = ed.Normal(loc=0., scale=1., name="loc")
x = ed.Normal(loc=loc, scale=0.5, sample_shape=5, name="x")
return x
def true_log_joint(loc, x):
log_prob = tf.reduce_sum(
tfd.Normal(loc=0., scale=1.).log_prob(loc))
log_prob += tf.reduce_sum(
tfd.Normal(loc=loc, scale=0.5).log_prob(x))
return log_prob
loc_value = 0.3
x_value = tf.random.normal([5])
log_joint = ed.make_log_joint_fn(normal_with_unknown_mean)
actual_log_prob = true_log_joint(loc_value, x_value)
expected_log_prob = log_joint(
loc=loc_value,
x=x_value,
f="https://github.com/tensorflow/probability/issues/160")
with self.assertRaises(LookupError):
_ = log_joint(loc=loc_value)
actual_log_prob_, expected_log_prob_ = self.evaluate(
[actual_log_prob, expected_log_prob])
self.assertEqual(actual_log_prob_, expected_log_prob_)
def testMakeLogJointFnTemplate(self):
"""Test `make_log_joint_fn` on program returned by tf1.make_template."""
def variational():
loc = tf1.get_variable("loc", [])
qz = ed.Normal(loc=loc, scale=0.5, name="qz")
return qz
def true_log_joint(loc, qz):
log_prob = tf.reduce_sum(
tfd.Normal(loc=loc, scale=0.5).log_prob(qz))
return log_prob
qz_value = 1.23
variational_template = tf1.make_template("variational", variational)
log_joint = ed.make_log_joint_fn(variational_template)
expected_log_prob = log_joint(qz=qz_value)
loc = tf1.trainable_variables("variational")[0]
actual_log_prob = true_log_joint(loc, qz_value)
with self.cached_session() as sess:
sess.run(tf1.initialize_all_variables())
actual_log_prob_, expected_log_prob_ = sess.run(
[actual_log_prob, expected_log_prob])
self.assertEqual(actual_log_prob_, expected_log_prob_)
def testMakeLogJointFnDynamic(self):
"""Test `make_log_joint_fn` on Edward program with stochastic control flow.
This verifies that Edward's program transformation is done by tracing the
execution at runtime (and not purely by static analysis). In particular,
the execution is controlled by random variable outcomes, which in turn is
controlled by the log-joint's inputs.
"""
def mixture_of_real_and_int():
loc = ed.Normal(loc=0., scale=1., name="loc")
flip = ed.Bernoulli(probs=0.5, name="flip")
if tf.equal(flip, 1):
x = ed.Normal(loc=loc, scale=0.5, sample_shape=5, name="x")
else:
x = ed.Poisson(rate=tf.nn.softplus(loc),
sample_shape=3,
name="x")
return x
def true_log_joint(loc, flip, x):
log_prob = tf.reduce_sum(
tfd.Normal(loc=0., scale=1.).log_prob(loc))
log_prob += tf.reduce_sum(tfd.Bernoulli(probs=0.5).log_prob(flip))
if tf.equal(flip, 1):
log_prob += tf.reduce_sum(
tfd.Normal(loc=loc, scale=0.5).log_prob(x))
else:
log_prob += tf.reduce_sum(
tfd.Poisson(rate=tf.nn.softplus(loc)).log_prob(x))
return log_prob
loc_value = 0.3
flip_value = tf.constant(1)
x_value = tf.random.normal([5])
log_joint = ed.make_log_joint_fn(mixture_of_real_and_int)
actual_log_prob = true_log_joint(loc_value, flip_value, x_value)
expected_log_prob = log_joint(loc=loc_value,
flip=flip_value,
x=x_value)
self.assertEqual(actual_log_prob, expected_log_prob)
loc_value = 1.2
flip_value = tf.constant(0)
x_value = tf.random.normal([3])
actual_log_prob = true_log_joint(loc_value, flip_value, x_value)
expected_log_prob = log_joint(loc=loc_value,
flip=flip_value,
x=x_value)
self.assertEqual(actual_log_prob, expected_log_prob)
def testMakeLogJointFnError(self):
"""Test `make_log_joint_fn` raises errors when `name`(s) not supplied."""
def normal_with_unknown_mean():
loc = ed.Normal(loc=0., scale=1., name="loc")
x = ed.Normal(loc=loc, scale=0.5, sample_shape=5)
return x
loc_value = 0.3
x_value = tf.random.normal([5])
log_joint = ed.make_log_joint_fn(normal_with_unknown_mean)
with self.assertRaises(KeyError):
_ = log_joint(loc=loc_value, x=x_value)
def testMakeLogJointFnConditional(self):
"""Test `make_log_joint_fn` on conditional Edward program."""
def linear_regression(features, prior_precision):
w = ed.Normal(loc=0.,
scale=tf.math.rsqrt(prior_precision),
sample_shape=features.shape[1],
name="w")
y = ed.Normal(loc=tf.tensordot(features, w, [[1], [0]]),
scale=1.,
name="y")
return y
features = tf.random.normal([3, 2])
prior_precision = 0.5
w_value = tf.random.normal([2])
y_value = tf.random.normal([3])
def true_log_joint(features, prior_precision, w, y):
log_prob = tf.reduce_sum(
tfd.Normal(loc=0.,
scale=tf.math.rsqrt(prior_precision)).log_prob(w))
log_prob += tf.reduce_sum(
tfd.Normal(loc=tf.tensordot(features, w, [[1], [0]]),
scale=1.).log_prob(y))
return log_prob
log_joint = ed.make_log_joint_fn(linear_regression)
actual_log_prob = true_log_joint(features, prior_precision, w_value,
y_value)
expected_log_prob = log_joint(features,
prior_precision,
y=y_value,
w=w_value)
with self.assertRaises(LookupError):
_ = log_joint(features, prior_precision, w=w_value)
actual_log_prob_, expected_log_prob_ = self.evaluate(
[actual_log_prob, expected_log_prob])
self.assertEqual(actual_log_prob_, expected_log_prob_)
def main(argv):
del argv # unused
if tf.io.gfile.isdir(FLAGS.model_dir):
tf.logging.warning(
"Warning: deleting old log directory at {}".format(FLAGS.model_dir))
tf.io.gfile.rmtree(FLAGS.model_dir)
tf.io.gfile.makedirs(FLAGS.model_dir)
tf.enable_v2_behavior()
print("GPU(s) available", tf.test.is_gpu_available())
if FLAGS.fake_data:
features = tf.random.normal([20, 55])
labels = tf.random.uniform([20], minval=0, maxval=2, dtype=tf.int32)
else:
features, labels = covertype()
print("Data set size", features.shape[0])
print("Number of features", features.shape[1])
log_joint = ed.make_log_joint_fn(logistic_regression)
def target_log_prob_fn(coeffs):
return log_joint(features=features, coeffs=coeffs, labels=labels)
# Initialize using a sample from 20 steps of NUTS. It is roughly a MAP
# estimate and is written explicitly to avoid differences in warm-starts
# between different implementations (e.g., Stan, PyMC3).
coeffs = tf.constant(
[+2.03420663e+00, -3.53567265e-02, -1.49223924e-01, -3.07049364e-01,
-1.00028366e-01, -1.46827862e-01, -1.64167881e-01, -4.20344204e-01,
+9.47479829e-02, -1.12681836e-02, +2.64442056e-01, -1.22087866e-01,
-6.00568838e-02, -3.79419506e-01, -1.06668741e-01, -2.97053963e-01,
-2.05253899e-01, -4.69537191e-02, -2.78072730e-02, -1.43250525e-01,
-6.77954629e-02, -4.34899796e-03, +5.90927452e-02, +7.23133609e-02,
+1.38526391e-02, -1.24497898e-01, -1.50733739e-02, -2.68872194e-02,
-1.80925727e-02, +3.47936489e-02, +4.03552800e-02, -9.98773426e-03,
+6.20188080e-02, +1.15002751e-01, +1.32145107e-01, +2.69109547e-01,
+2.45785132e-01, +1.19035013e-01, -2.59744357e-02, +9.94279515e-04,
+3.39266285e-02, -1.44057125e-02, -6.95222765e-02, -7.52013028e-02,
+1.21171586e-01, +2.29205526e-02, +1.47308692e-01, -8.34354162e-02,
-9.34122875e-02, -2.97472421e-02, -3.03937674e-01, -1.70958012e-01,
-1.59496680e-01, -1.88516974e-01, -1.20889175e+00])
# Initialize step size via result of 50 warmup steps from Stan.
step_size = 0.00167132
kernel = profile(no_u_turn_sampler.kernel)
coeffs_samples = []
target_log_prob = None
grads_target_log_prob = None
for step in range(FLAGS.max_steps):
print("Step", step)
[
[coeffs],
target_log_prob,
grads_target_log_prob,
] = kernel(target_log_prob_fn=target_log_prob_fn,
current_state=[coeffs],
step_size=[step_size],
seed=step,
current_target_log_prob=target_log_prob,
current_grads_target_log_prob=grads_target_log_prob)
coeffs_samples.append(coeffs)
for coeffs_sample in coeffs_samples:
plt.plot(coeffs_sample.numpy())
filename = os.path.join(FLAGS.model_dir, "coeffs_samples.png")
plt.savefig(filename)
print("Figure saved as", filename)
plt.close()
