dtype issues with tfp.mcmc.sample_chain

I apologize in advance if this is just some misunderstanding on my part and not a bug. I am trying to run MCMC to estimate the logit vector that parameterizes a binomial distribution modeling 5000 outcomes of rolling a 6-sided die (each represented by a one-hot encoded vectors). However, I am running into issues with dtype clashes somewhere. The following code produces the following error output (tf.version='2.1.0', tfp.version='0.8.0-rc0'):

import tensorflow as tf
import tensorflow_probability as tfp
import numpy as np
tfd = tfp.distributions
tfb = tfp.bijectors

tf.config.list_physical_devices('GPU')

# Get the roll data, the roll sums (x_k's), and the number of rolls (n)
roll_data = tf.constant(np.loadtxt('unfair_dice.txt'), dtype=tf.float64)
roll_sums = tf.reduce_sum(roll_data, axis=0)
num_rolls = roll_data.shape[0]

        

# Some kind of problem when checking the accept probability
def unnormalized_log_prob(data, logit_p_vec_est):
    dist_est = tfd.Multinomial(total_count=data.shape[0], logits=logit_p_vec_est)
    class_counts = class_counts = tf.reduce_sum(data, axis=0)
    
    return tf.constant(tf.reduce_mean(dist_est.log_prob(class_counts)), dtype=tf.float64)

hmc_unnormalized_log_prob = lambda p_vec: unnormalized_log_prob(roll_data, p_vec)

initial_chain_state = [tf.constant(np.ones(1), dtype=tf.float64)]
step_size = tf.constant(0.01, dtype=tf.float64)
leapfrog_steps_num = 3
sample_size = int(10e3)
burnin = int(10e3)

adaptive_hmc = tfp.mcmc.SimpleStepSizeAdaptation(
    tfp.mcmc.HamiltonianMonteCarlo(
        target_log_prob_fn=hmc_unnormalized_log_prob,
        num_leapfrog_steps=leapfrog_steps_num,
        step_size=step_size),
    num_adaptation_steps=int(burnin*0.8))



samples_hmc = tfp.mcmc.sample_chain(
        num_results=sample_size,
        num_burnin_steps=burnin,
        current_state=initial_chain_state,
        kernel=adaptive_hmc,
        trace_fn=None)

---------------------------------------------------------------------------
InvalidArgumentError                      Traceback (most recent call last)
<ipython-input-22-42e38bc23afb> in <module>
     43         current_state=initial_chain_state,
     44         kernel=adaptive_hmc,
---> 45         trace_fn=None)

~/anaconda3/envs/am216/lib/python3.7/site-packages/tensorflow_probability/python/mcmc/sample.py in sample_chain(num_results, current_state, previous_kernel_results, kernel, num_burnin_steps, num_steps_between_results, trace_fn, return_final_kernel_results, parallel_iterations, name)
    360                                             trace_fn(*state_and_results)),
    361         # pylint: enable=g-long-lambda
--> 362         parallel_iterations=parallel_iterations)
    363 
    364     if return_final_kernel_results:

~/anaconda3/envs/am216/lib/python3.7/site-packages/tensorflow_probability/python/mcmc/internal/util.py in trace_scan(loop_fn, initial_state, elems, trace_fn, parallel_iterations, name)
    369         body=_body,
    370         loop_vars=(0, initial_state, trace_arrays),
--> 371         parallel_iterations=parallel_iterations)
    372 
    373     stacked_trace = tf.nest.map_structure(lambda x: x.stack(), trace_arrays)

~/anaconda3/envs/am216/lib/python3.7/site-packages/tensorflow_core/python/ops/control_flow_ops.py in while_loop_v2(cond, body, loop_vars, shape_invariants, parallel_iterations, back_prop, swap_memory, maximum_iterations, name)
   2476       name=name,
   2477       maximum_iterations=maximum_iterations,
-> 2478       return_same_structure=True)
   2479 
   2480 

~/anaconda3/envs/am216/lib/python3.7/site-packages/tensorflow_core/python/ops/control_flow_ops.py in while_loop(cond, body, loop_vars, shape_invariants, parallel_iterations, back_prop, swap_memory, name, maximum_iterations, return_same_structure)
   2712                                               list(loop_vars))
   2713       while cond(*loop_vars):
-> 2714         loop_vars = body(*loop_vars)
   2715         if try_to_pack and not isinstance(loop_vars, (list, _basetuple)):
   2716           packed = True

~/anaconda3/envs/am216/lib/python3.7/site-packages/tensorflow_probability/python/mcmc/internal/util.py in _body(i, state, trace_arrays)
    358 
    359     def _body(i, state, trace_arrays):
--> 360       state = loop_fn(state, elems_array.read(i))
    361       trace_arrays = tf.nest.pack_sequence_as(trace_arrays, [
    362           a.write(i, v) for a, v in zip(

~/anaconda3/envs/am216/lib/python3.7/site-packages/tensorflow_probability/python/mcmc/sample.py in _trace_scan_fn(state_and_results, num_steps)
    344           body_fn=kernel.one_step,
    345           initial_loop_vars=list(state_and_results),
--> 346           parallel_iterations=parallel_iterations)
    347       return next_state, current_kernel_results
    348 

~/anaconda3/envs/am216/lib/python3.7/site-packages/tensorflow_probability/python/mcmc/internal/util.py in smart_for_loop(loop_num_iter, body_fn, initial_loop_vars, parallel_iterations, name)
    285           body=lambda i, *args: [i + 1] + list(body_fn(*args)),
    286           loop_vars=[np.int32(0)] + initial_loop_vars,
--> 287           parallel_iterations=parallel_iterations
    288       )[1:]
    289     result = initial_loop_vars

~/anaconda3/envs/am216/lib/python3.7/site-packages/tensorflow_core/python/ops/control_flow_ops.py in while_loop_v2(cond, body, loop_vars, shape_invariants, parallel_iterations, back_prop, swap_memory, maximum_iterations, name)
   2476       name=name,
   2477       maximum_iterations=maximum_iterations,
-> 2478       return_same_structure=True)
   2479 
   2480 

~/anaconda3/envs/am216/lib/python3.7/site-packages/tensorflow_core/python/ops/control_flow_ops.py in while_loop(cond, body, loop_vars, shape_invariants, parallel_iterations, back_prop, swap_memory, name, maximum_iterations, return_same_structure)
   2712                                               list(loop_vars))
   2713       while cond(*loop_vars):
-> 2714         loop_vars = body(*loop_vars)
   2715         if try_to_pack and not isinstance(loop_vars, (list, _basetuple)):
   2716           packed = True

~/anaconda3/envs/am216/lib/python3.7/site-packages/tensorflow_probability/python/mcmc/internal/util.py in <lambda>(i, *args)
    283       return tf.while_loop(
    284           cond=lambda i, *args: i < loop_num_iter,
--> 285           body=lambda i, *args: [i + 1] + list(body_fn(*args)),
    286           loop_vars=[np.int32(0)] + initial_loop_vars,
    287           parallel_iterations=parallel_iterations

~/anaconda3/envs/am216/lib/python3.7/site-packages/tensorflow_probability/python/mcmc/simple_step_size_adaptation.py in one_step(self, current_state, previous_kernel_results)
    388 
    389         new_step_size_part = mcmc_util.choose(
--> 390             reduced_log_accept_prob > log_target_accept_prob,
    391             step_size_part * (1. + previous_kernel_results.adaptation_rate),
    392             step_size_part / (1. + previous_kernel_results.adaptation_rate))

~/anaconda3/envs/am216/lib/python3.7/site-packages/tensorflow_core/python/ops/gen_math_ops.py in greater(x, y, name)
   3989         raise
   3990     except _core._NotOkStatusException as e:
-> 3991       _ops.raise_from_not_ok_status(e, name)
   3992   # Add nodes to the TensorFlow graph.
   3993   try:

~/anaconda3/envs/am216/lib/python3.7/site-packages/tensorflow_core/python/framework/ops.py in raise_from_not_ok_status(e, name)
   6604   message = e.message + (" name: " + name if name is not None else "")
   6605   # pylint: disable=protected-access
-> 6606   six.raise_from(core._status_to_exception(e.code, message), None)
   6607   # pylint: enable=protected-access
   6608 

~/anaconda3/envs/am216/lib/python3.7/site-packages/six.py in raise_from(value, from_value)

InvalidArgumentError: cannot compute Greater as input #1(zero-based) was expected to be a double tensor but is a float tensor [Op:Greater]

However, changing all dtypes to tf.float32 makes the code run smoothly and gives the correct logit estimates. Does this mean that one of the tfp.mcmc calls I am using is only compatable with tf.float32? If so, is there somewhere in the documentation that indicates this? If not, could this be a bug somwhere in the tfp code? Thanks in advance.

Reproduce the dtype mismatch using the provided float64 example, focusing first on tfp.mcmc.sample_chain and SimpleStepSizeAdaptation. Compare the float64 and float32 paths and inspect the reported Greater operation; done means the issue is either fixed for float64 or its supported dtype behavior is documented with a regression test.