Increase shard counts for ode_test and matern_test to avoid timeouts.

PiperOrigin-RevId: 456611097

Author: emilyfertig

tensorflow_probability/python/math/ode/BUILD

@@ -131,7 +131,7 @@ multi_substrate_py_test(
     name = "ode_test",
     size = "large",
     srcs = ["ode_test.py"],
-    shard_count = 8,
+    shard_count = 10,
     deps = [
         # absl/testing:parameterized dep,
         # numpy dep,

tensorflow_probability/python/math/psd_kernels/BUILD

@@ -182,7 +182,7 @@ multi_substrate_py_test(
     srcs = ["matern_test.py"],
     jax_size = "medium",
     numpy_tags = ["notap"],
-    shard_count = 2,
+    shard_count = 4,
     deps = [
         # absl/testing:parameterized dep,
         # numpy dep,