diff --git a/tensorflow_probability/python/layers/BUILD b/tensorflow_probability/python/layers/BUILD
index 85150560cf..5ae42b98cf 100644
--- a/tensorflow_probability/python/layers/BUILD
+++ b/tensorflow_probability/python/layers/BUILD
@@ -138,6 +138,7 @@ py_test(
     name = "distribution_layer_test",
     size = "large",
     srcs = ["distribution_layer_test.py"],
+    shard_count = 5,
     deps = [
         # numpy dep,
         # tensorflow dep,
diff --git a/tensorflow_probability/python/layers/distribution_layer_test.py b/tensorflow_probability/python/layers/distribution_layer_test.py
index ac5b2a6e67..7408ccb960 100644
--- a/tensorflow_probability/python/layers/distribution_layer_test.py
+++ b/tensorflow_probability/python/layers/distribution_layer_test.py
@@ -80,6 +80,7 @@ def setUp(self):
         self.train_size, *self.input_shape) > 0.75).astype(np.float32)
     self.x_test = (np.random.rand(
         self.test_size, *self.input_shape) > 0.75).astype(np.float32)
+    super(EndToEndTest, self).setUp()
 
   def test_keras_sequential_api(self):
     """Test `DistributionLambda`s are composable via Keras `Sequential` API."""