Enable JointDistribution tests, even though they don't currently pass…

… (CI is disabled). The change to rewrite.py also enables importing/using the JD classes in np/jax tests like cholesky_lkj w/ the logdet fix.

PiperOrigin-RevId: 288711875

tensorflow_probability/python/distributions/BUILD

@@ -2037,9 +2037,11 @@ multi_substrate_py_test(
     ],
 )
 
-py_test(
+multi_substrate_py_test(
     name = "joint_distribution_coroutine_test",
     srcs = ["joint_distribution_coroutine_test.py"],
+    jax_tags = ["notap"],
+    numpy_tags = ["notap"],
     deps = [
         # numpy dep,
         # tensorflow dep,
@@ -2048,9 +2050,11 @@ py_test(
     ],
 )
 
-py_test(
+multi_substrate_py_test(
     name = "joint_distribution_named_test",
     srcs = ["joint_distribution_named_test.py"],
+    jax_tags = ["notap"],
+    numpy_tags = ["notap"],
     deps = [
         # absl/testing:parameterized dep,
         # tensorflow dep,
@@ -2059,9 +2063,11 @@ py_test(
     ],
 )
 
-py_test(
+multi_substrate_py_test(
     name = "joint_distribution_sequential_test",
     srcs = ["joint_distribution_sequential_test.py"],
+    jax_tags = ["notap"],
+    numpy_tags = ["notap"],
     deps = [
         # absl/testing:parameterized dep,
         # numpy dep,

tensorflow_probability/python/experimental/substrates/meta/rewrite.py

@@ -47,9 +47,7 @@
     'bijectors':
         ('masked_autoregressive', 'scale_matvec_lu', 'real_nvp'),
     'distributions':
-        ('joint_distribution', 'joint_distribution_coroutine',
-         'joint_distribution_named', 'joint_distribution_sequential',
-         'internal.moving_stats'),
+        ('internal.moving_stats',),
     'math':
         ('ode', 'diag_jacobian',
          'interpolation', 'minimize', 'root_search', 'sparse'),