From 8d1400fe608b782addc06428cbb330267d8fde00 Mon Sep 17 00:00:00 2001
From: Alexey Radul <axch@google.com>
Date: Mon, 12 Aug 2019 15:00:23 -0700
Subject: [PATCH] Shard distribution_layer_test so it doesn't time out as
 often.

PiperOrigin-RevId: 263012414
---
 tensorflow_probability/python/layers/BUILD                      | 1 +
 tensorflow_probability/python/layers/distribution_layer_test.py | 1 +
 2 files changed, 2 insertions(+)

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."""