diff --git a/tensor2tensor/layers/bayes.py b/tensor2tensor/layers/bayes.py
index 76e6ae5f8..1b60cd7b1 100644
--- a/tensor2tensor/layers/bayes.py
+++ b/tensor2tensor/layers/bayes.py
@@ -19,8 +19,11 @@
 from __future__ import division
 from __future__ import print_function
 
+import functools
 import tensorflow as tf
 
+from tensorflow_probability import edward2 as ed
+
 
 class Softplus(tf.keras.constraints.Constraint):
   """Softplus constraint."""
@@ -110,11 +113,7 @@ def __call__(self, shape=None, dtype=None, partition_info=None):
       raise ValueError('A TrainableInitializer must be built by a layer before '
                        'usage, and is currently only compatible with Bayesian '
                        'layers.')
-    noise = tf.random_normal(self.shape, dtype=self.dtype, seed=self.seed)
-    output = self.mean + self.stddev * noise
-    # TODO(trandustin): Hack to store parameters so KL reg. can operate on them.
-    output._parameters = (self.mean, self.stddev)  # pylint: disable=protected-access
-    return output
+    return ed.Normal(loc=self.mean, scale=self.stddev)
 
   def get_config(self):
     return {
@@ -151,12 +150,11 @@ def __init__(self, mean=0., stddev=1.):
     self.stddev = stddev
 
   def __call__(self, x):
-    mean, stddev = x._parameters  # pylint: disable=protected-access
-    variance2 = tf.square(self.stddev)
-    variance_ratio = tf.square(stddev) / variance2
-    regularization = tf.square(mean - self.mean) / (2. * variance2)
-    regularization += (variance_ratio - 1. - tf.log(variance_ratio)) / 2.
-    return regularization
+    """Computes regularization given an ed.Normal random variable as input."""
+    if not isinstance(x, ed.RandomVariable):
+      raise ValueError('Input must be an ed.RandomVariable.')
+    random_variable = ed.Normal(loc=self.mean, scale=self.stddev)
+    return random_variable.distribution.kl_divergence(x.distribution)
 
   def get_config(self):
     return {
@@ -276,3 +274,18 @@ def build(self, input_shape):
     else:
       self.bias = None
     self.built = True
+
+  # TODO(trandustin): Waiting on T2T to drop dependence on
+  # TF<=1.12rc2. A TF commit enables tf.colocate_with to work for
+  # Tensor-like inputs. This lets us use the parent method instead of
+  # this one.
+  def _handle_weight_regularization(self, name, variable, regularizer):
+    """Create lambdas which compute regularization losses."""
+
+    def _loss_for_variable(v):
+      """Creates a regularization loss `Tensor` for variable `v`."""
+      with tf.name_scope(name + '/Regularizer'):
+        regularization = regularizer(v)
+      return regularization
+
+    self.add_loss(functools.partial(_loss_for_variable, variable))
diff --git a/tensor2tensor/layers/bayes_test.py b/tensor2tensor/layers/bayes_test.py
index c9e754438..521d7de06 100644
--- a/tensor2tensor/layers/bayes_test.py
+++ b/tensor2tensor/layers/bayes_test.py
@@ -29,25 +29,20 @@
 
 class BayesTest(parameterized.TestCase, tf.test.TestCase):
 
-  # TODO(trandustin): Remove the hack in the code, or re-enable once T2T drops
-  # support for TF 1.10
-  # @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
   def testDenseReparameterizationKernel(self):
     inputs = tf.to_float(np.random.rand(5, 3, 12))
     layer = bayes.DenseReparameterization(4, activation=tf.nn.relu)
     outputs1 = layer(inputs)
     outputs2 = layer(inputs)
     self.evaluate(tf.global_variables_initializer())
-    # res1, res2 = self.evaluate([outputs1, outputs2])
-    res1, _ = self.evaluate([outputs1, outputs2])
+    res1, res2 = self.evaluate([outputs1, outputs2])
     self.assertEqual(res1.shape, (5, 3, 4))
     self.assertAllGreaterEqual(res1, 0.)
-    # self.assertNotAllClose(res1, res2)
+    self.assertNotAllClose(res1, res2)
     layer.get_config()
 
-  # TODO(trandustin): Remove the hack in the code, or re-enable once T2T drops
-  # support for TF 1.10
-  # @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
   def testDenseReparameterizationBias(self):
     inputs = tf.to_float(np.random.rand(5, 3, 12))
     layer = bayes.DenseReparameterization(4, kernel_initializer="zero",
@@ -56,15 +51,12 @@ def testDenseReparameterizationBias(self):
     outputs1 = layer(inputs)
     outputs2 = layer(inputs)
     self.evaluate(tf.global_variables_initializer())
-    # res1, res2 = self.evaluate([outputs1, outputs2])
-    res1, _ = self.evaluate([outputs1, outputs2])
+    res1, res2 = self.evaluate([outputs1, outputs2])
     self.assertEqual(res1.shape, (5, 3, 4))
     self.assertAllGreaterEqual(res1, 0.)
-    # self.assertNotAllClose(res1, res2)
+    self.assertNotAllClose(res1, res2)
 
-  # TODO(trandustin): Remove the hack in the code, or re-enable once T2T drops
-  # support for TF 1.10
-  # @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
   def testDenseReparameterizationDeterministic(self):
     inputs = tf.to_float(np.random.rand(5, 3, 12))
     layer = bayes.DenseReparameterization(4, kernel_initializer="zero",
@@ -73,15 +65,12 @@ def testDenseReparameterizationDeterministic(self):
     outputs1 = layer(inputs)
     outputs2 = layer(inputs)
     self.evaluate(tf.global_variables_initializer())
-    # res1, res2 = self.evaluate([outputs1, outputs2])
-    res1, _ = self.evaluate([outputs1, outputs2])
+    res1, res2 = self.evaluate([outputs1, outputs2])
     self.assertEqual(res1.shape, (5, 3, 4))
     self.assertAllGreaterEqual(res1, 0.)
-    # self.assertAllClose(res1, res2)
+    self.assertAllClose(res1, res2)
 
-  # TODO(trandustin): Remove the hack in the code, or re-enable once T2T drops
-  # support for TF 1.10
-  # @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
   def testDenseReparameterizationModel(self):
     inputs = tf.to_float(np.random.rand(3, 4, 4, 1))
     model = tf.keras.Sequential([