Restore tfd.Mixture to jax substrate.

PiperOrigin-RevId: 455459945

Author: Googler

tensorflow_probability/python/distributions/BUILD

@@ -36,7 +36,6 @@ multi_substrate_py_library(
     srcs = ["__init__.py"],
     substrates_omit_deps = [
         ":pixel_cnn",
-        ":mixture",
     ],
     deps = [
         ":autoregressive",
@@ -3322,7 +3321,6 @@ multi_substrate_py_test(
     name = "mixture_test",
     size = "medium",
     srcs = ["mixture_test.py"],
-    jax_tags = ["notap"],
     numpy_tags = ["notap"],
     deps = [
         # hypothesis dep,

tensorflow_probability/python/distributions/mixture.py

@@ -347,7 +347,7 @@ def _sample_n(self, n, seed=None):
     mask = distribution_util.pad_mixture_dimensions(
         mask, self, self._cat,
         tensorshape_util.rank(self._static_event_shape))  # [n, B, k, [1]*e]
-    if x.dtype.is_floating:
+    if dtype_util.is_floating(x.dtype):
       masked = tf.math.multiply_no_nan(x, mask)
     else:
       masked = x * mask