Remove usage of deprecated layer API.

PiperOrigin-RevId: 427561526

Author: fchollet

tensorflow_probability/python/layers/conv_variational_test.py

@@ -245,14 +245,14 @@ def _testKLPenaltyKernel(self, layer_class):  # pylint: disable=invalid-name
       inputs = self.maybe_transpose_tensor(inputs)
 
       # No keys.
-      input_dependent_losses = layer.get_losses_for(inputs=None)
+      input_dependent_losses = layer.losses
       self.assertEqual(len(layer.losses), 0)
       self.assertListEqual(layer.losses, input_dependent_losses)
 
       _ = layer(inputs)
 
       # Yes keys.
-      input_dependent_losses = layer.get_losses_for(inputs=None)
+      input_dependent_losses = layer.losses
       self.assertEqual(len(layer.losses), 1)
       self.assertEqual(layer.losses[0].shape, ())
       self.assertListEqual(layer.losses, input_dependent_losses)
@@ -277,18 +277,14 @@ def _testKLPenaltyBoth(self, layer_class):  # pylint: disable=invalid-name
       inputs = self.maybe_transpose_tensor(inputs)
 
       # No keys.
-      input_dependent_losses = layer.get_losses_for(inputs=None)
       self.assertEqual(len(layer.losses), 0)
-      self.assertListEqual(layer.losses, input_dependent_losses)
 
       _ = layer(inputs)
 
       # Yes keys.
-      input_dependent_losses = layer.get_losses_for(inputs=None)
       self.assertEqual(len(layer.losses), 2)
       self.assertEqual(layer.losses[0].shape, ())
       self.assertEqual(layer.losses[1].shape, ())
-      self.assertListEqual(layer.losses, input_dependent_losses)
 
   def _testConvSetUp(self, layer_class, batch_size, depth=None,
                      height=None, width=None, channels=None, filters=None,
@@ -350,7 +346,7 @@ def _testConvSetUp(self, layer_class, batch_size, depth=None,
 
     outputs = layer(inputs)
 
-    kl_penalty = layer.get_losses_for(inputs=None)
+    kl_penalty = layer.losses
     return (kernel_posterior, kernel_prior, kernel_divergence,
             bias_posterior, bias_prior, bias_divergence,
             layer, inputs, outputs, kl_penalty, kernel_shape)

tensorflow_probability/python/layers/dense_variational_test.py

@@ -140,17 +140,13 @@ def _testKLPenaltyKernel(self, layer_class):
       inputs = tf.random.uniform([2, 3], seed=1)
 
       # No keys.
-      input_dependent_losses = layer.get_losses_for(inputs=None)
       self.assertEqual(len(layer.losses), 0)
-      self.assertListEqual(layer.losses, input_dependent_losses)
 
       _ = layer(inputs)
 
       # Yes keys.
-      input_dependent_losses = layer.get_losses_for(inputs=None)
       self.assertEqual(len(layer.losses), 1)
       self.assertEqual(layer.losses[0].shape, ())
-      self.assertListEqual(layer.losses, input_dependent_losses)
 
   def _testKLPenaltyBoth(self, layer_class):
     with self.cached_session():
@@ -161,18 +157,14 @@ def _testKLPenaltyBoth(self, layer_class):
       inputs = tf.random.uniform([2, 3], seed=1)
 
       # No keys.
-      input_dependent_losses = layer.get_losses_for(inputs=None)
       self.assertEqual(len(layer.losses), 0)
-      self.assertListEqual(layer.losses, input_dependent_losses)
 
       _ = layer(inputs)
 
       # Yes keys.
-      input_dependent_losses = layer.get_losses_for(inputs=None)
       self.assertEqual(len(layer.losses), 2)
       self.assertEqual(layer.losses[0].shape, ())
       self.assertEqual(layer.losses[1].shape, ())
-      self.assertListEqual(layer.losses, input_dependent_losses)
 
   def _testDenseSetUp(self, layer_class, batch_size, in_size, out_size,
                       **kwargs):
@@ -215,7 +207,7 @@ def _testDenseSetUp(self, layer_class, batch_size, in_size, out_size,
 
     outputs = layer(inputs)
 
-    kl_penalty = layer.get_losses_for(inputs=None)
+    kl_penalty = layer.losses
     return (kernel_posterior, kernel_prior, kernel_divergence,
             bias_posterior, bias_prior, bias_divergence,
             layer, inputs, outputs, kl_penalty)