Cast x to logits.dtype in onehot_categorical.

Otherwise, dist.log_prob(dist.sample()) has an error in reduce_logsumexp inside _assert_valid_sample when validate_args=True (TypeError: Value passed to parameter 'x' has DataType int32 not in list of allowed values: bfloat16, float16, float32, float64).

PiperOrigin-RevId: 239992007

tensorflow_probability/python/distributions/onehot_categorical.py

@@ -181,6 +181,7 @@ def _sample_n(self, n, seed=None):
     return ret
 
   def _log_prob(self, x):
+    x = tf.cast(x, self.logits.dtype)
     x = self._assert_valid_sample(x)
     # broadcast logits or x if need be.
     logits = self.logits
@@ -227,7 +228,7 @@ def _assert_valid_sample(self, x):
     return distribution_util.with_dependencies([
         tf.compat.v1.assert_non_positive(x),
         tf.compat.v1.assert_near(
-            tf.zeros([], dtype=self.dtype),
+            tf.zeros([], dtype=self.logits.dtype),
             tf.reduce_logsumexp(input_tensor=x, axis=[-1])),
     ], x)