Merge pull request #907 from emilyfertig/r0.10

For R0.10-rc1, with cherry-pick to fix convolutional layers.

Author: emilyfertig

tensorflow_probability/python/layers/conv_variational.py

@@ -197,7 +197,6 @@ def build(self, input_shape):
       self.kernel_prior = self.kernel_prior_fn(
           dtype, kernel_shape, 'kernel_prior',
           self.trainable, self.add_variable)
-    self._built_kernel_divergence = False
 
     if self.bias_posterior_fn is None:
       self.bias_posterior = None
@@ -212,7 +211,6 @@ def build(self, input_shape):
       self.bias_prior = self.bias_prior_fn(
           dtype, (self.filters,), 'bias_prior',
           self.trainable, self.add_variable)
-    self._built_bias_divergence = False
 
     self.input_spec = tf.keras.layers.InputSpec(
         ndim=self.rank + 2, axes={channel_axis: input_dim})
@@ -234,20 +232,16 @@ def call(self, inputs):
     outputs = self._apply_variational_bias(outputs)
     if self.activation is not None:
       outputs = self.activation(outputs)
-    if not self._built_kernel_divergence:
-      self._apply_divergence(self.kernel_divergence_fn,
-                             self.kernel_posterior,
-                             self.kernel_prior,
-                             self.kernel_posterior_tensor,
-                             name='divergence_kernel')
-      self._built_kernel_divergence = True
-    if not self._built_bias_divergence:
-      self._apply_divergence(self.bias_divergence_fn,
-                             self.bias_posterior,
-                             self.bias_prior,
-                             self.bias_posterior_tensor,
-                             name='divergence_bias')
-      self._built_bias_divergence = True
+    self._apply_divergence(self.kernel_divergence_fn,
+                           self.kernel_posterior,
+                           self.kernel_prior,
+                           self.kernel_posterior_tensor,
+                           name='divergence_kernel')
+    self._apply_divergence(self.bias_divergence_fn,
+                           self.bias_posterior,
+                           self.bias_prior,
+                           self.bias_posterior_tensor,
+                           name='divergence_bias')
     return outputs
 
   def compute_output_shape(self, input_shape):

tensorflow_probability/python/layers/conv_variational_test.py

@@ -21,7 +21,6 @@
 # Dependency imports
 import numpy as np
 
-import tensorflow.compat.v1 as tf1
 import tensorflow.compat.v2 as tf
 import tensorflow_probability as tfp
 
@@ -180,12 +179,12 @@ def __call__(self, *args, **kwargs):
 @test_util.test_all_tf_execution_regimes
 class ConvVariational(object):
 
-  def maybe_transpose_inputs(self, inputs):
+  def maybe_transpose_tensor(self, tensor):
     if self.data_format == 'channels_first':
-      order = channels_last_to_first(list(range(inputs.shape.rank)))
-      return tf.transpose(a=inputs, perm=order)
+      order = channels_last_to_first(list(range(tensor.shape.rank)))
+      return tf.transpose(a=tensor, perm=order)
     else:
-      return inputs
+      return tensor
 
   def _testKerasLayer(self, layer_class):  # pylint: disable=invalid-name
     def kernel_posterior_fn(dtype, shape, name, trainable, add_variable_fn):
@@ -241,7 +240,7 @@ def _testKLPenaltyKernel(self, layer_class):  # pylint: disable=invalid-name
       elif layer_class in (tfp.layers.Convolution3DReparameterization,
                            tfp.layers.Convolution3DFlipout):
         inputs = tf.random.uniform([2, 3, 3, 3, 1], seed=1)
-      inputs = self.maybe_transpose_inputs(inputs)
+      inputs = self.maybe_transpose_tensor(inputs)
 
       # No keys.
       input_dependent_losses = layer.get_losses_for(inputs=None)
@@ -273,7 +272,7 @@ def _testKLPenaltyBoth(self, layer_class):  # pylint: disable=invalid-name
       elif layer_class in (tfp.layers.Convolution3DReparameterization,
                            tfp.layers.Convolution3DFlipout):
         inputs = tf.random.uniform([2, 3, 3, 3, 1], seed=1)
-      inputs = self.maybe_transpose_inputs(inputs)
+      inputs = self.maybe_transpose_tensor(inputs)
 
       # No keys.
       input_dependent_losses = layer.get_losses_for(inputs=None)
@@ -307,7 +306,7 @@ def _testConvSetUp(self, layer_class, batch_size, depth=None,
       inputs = tf.random.uniform([batch_size, depth, height, width, channels],
                                  seed=seed())
       kernel_size = (2, 2, 2)
-    inputs = self.maybe_transpose_inputs(inputs)
+    inputs = self.maybe_transpose_tensor(inputs)
 
     kernel_shape = kernel_size + (channels, filters)
     kernel_posterior = MockDistribution(
@@ -547,7 +546,7 @@ def _testRandomConvFlipout(self, layer_class):  # pylint: disable=invalid-name
         inputs = tf.random.uniform([batch_size, depth, height, width, channels],
                                    seed=seed())
         kernel_size = (2, 2, 2)
-      inputs = self.maybe_transpose_inputs(inputs)
+      inputs = self.maybe_transpose_tensor(inputs)
 
       kernel_shape = kernel_size + (channels, filters)
       bias_size = (filters,)
@@ -597,26 +596,30 @@ def _testRandomConvFlipout(self, layer_class):  # pylint: disable=invalid-name
                       np.prod(outputs_one_.shape))
 
   def _testLayerInSequential(self, layer_class):  # pylint: disable=invalid-name
-    with self.cached_session() as sess:
-      if layer_class in (tfp.layers.Convolution1DReparameterization,
-                         tfp.layers.Convolution1DFlipout):
-        inputs = tf.random.uniform([2, 3, 1])
-      elif layer_class in (tfp.layers.Convolution2DReparameterization,
-                           tfp.layers.Convolution2DFlipout):
-        inputs = tf.random.uniform([2, 3, 3, 1])
-      elif layer_class in (tfp.layers.Convolution3DReparameterization,
-                           tfp.layers.Convolution3DFlipout):
-        inputs = tf.random.uniform([2, 3, 3, 3, 1])
-      inputs = self.maybe_transpose_inputs(inputs)
+    if layer_class in (tfp.layers.Convolution1DReparameterization,
+                       tfp.layers.Convolution1DFlipout):
+      inputs = tf.random.uniform([2, 3, 1])
+      outputs = tf.random.uniform([2, 1, 2])
+    elif layer_class in (tfp.layers.Convolution2DReparameterization,
+                         tfp.layers.Convolution2DFlipout):
+      inputs = tf.random.uniform([2, 3, 3, 1])
+      outputs = tf.random.uniform([2, 1, 1, 2])
+    elif layer_class in (tfp.layers.Convolution3DReparameterization,
+                         tfp.layers.Convolution3DFlipout):
+      inputs = tf.random.uniform([2, 3, 3, 3, 1])
+      outputs = tf.random.uniform([2, 1, 1, 1, 2])
+    inputs = self.maybe_transpose_tensor(inputs)
+    outputs = self.maybe_transpose_tensor(outputs)
+
+    net = tf.keras.Sequential([
+        layer_class(filters=2, kernel_size=3, data_format=self.data_format),
+        layer_class(filters=2, kernel_size=1, data_format=self.data_format)])
 
-      net = tf.keras.Sequential([
-          layer_class(filters=2, kernel_size=3, data_format=self.data_format),
-          layer_class(filters=2, kernel_size=1, data_format=self.data_format)])
-      output = net(inputs)
+    net.compile(loss='mse', optimizer='adam')
+    net.fit(inputs, outputs, batch_size=2, epochs=3, steps_per_epoch=2)
 
-      # Verify that the network runs without errors
-      sess.run(tf1.global_variables_initializer())
-      sess.run(output)
+    batch_output = self.evaluate(net(inputs))
+    self.assertAllEqual(outputs.shape, batch_output.shape)
 
   def testKerasLayerConvolution1DReparameterization(self):
     self._testKerasLayer(tfp.layers.Convolution1DReparameterization)

tensorflow_probability/python/version.py

@@ -24,7 +24,7 @@
 # stable release (indicated by `_VERSION_SUFFIX = ''`). Outside the context of a
 # release branch, the current version is by default assumed to be a
 # 'development' version, labeled 'dev'.
-_VERSION_SUFFIX = 'rc0'
+_VERSION_SUFFIX = 'rc1'
 
 # Example, '0.4.0-dev'
 __version__ = '.'.join([