Increase time limits for two tests, lest they time out sporadically

PiperOrigin-RevId: 292396433
csveda · Jan 30, 2020 · 540b9b1 · 540b9b1
1 parent b7fd244
commit 540b9b1
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 0 deletions.
diff --git a/tensorflow_probability/python/distributions/BUILD b/tensorflow_probability/python/distributions/BUILD
@@ -1734,6 +1734,7 @@ multi_substrate_py_test(
 multi_substrate_py_test(
     name = "dirichlet_test",
     srcs = ["dirichlet_test.py"],
+    jax_size = "medium",
     deps = [
         # numpy dep,
         # scipy dep,

diff --git a/tensorflow_probability/python/math/psd_kernels/BUILD b/tensorflow_probability/python/math/psd_kernels/BUILD
@@ -140,6 +140,7 @@ multi_substrate_py_test(
     name = "matern_test",
     size = "small",
     srcs = ["matern_test.py"],
+    jax_size = "medium",
     numpy_tags = ["notap"],
     deps = [
         # absl/testing:parameterized dep,