Sampling tfd.Mixture errors on GPU

The following

import tensorflow_probability as tfp
tfd = tfp.distributions
mix = tfd.Mixture(cat = tfd.Categorical([1, 1]), 
                  components = [tfd.Normal(0, 1), tfd.Normal(0, 2)])
with tf.device("/gpu:0"):
    mix.sample()

fails with the following error

NotFoundError                             Traceback (most recent call last)
<ipython-input-20-f996fdfe102a> in <module>
      4                   components = [tfd.Normal(0, 1), tfd.Normal(0, 2)])
      5 with tf.device("/gpu:0"):
----> 6     mix.sample()

~/tfnightly/lib/python3.6/site-packages/tensorflow_probability/python/distributions/distribution.py in sample(self, sample_shape, seed, name, **kwargs)
    838       samples: a `Tensor` with prepended dimensions `sample_shape`.
    839     """
--> 840     return self._call_sample_n(sample_shape, seed, name, **kwargs)
    841 
    842   def _call_log_prob(self, value, name, **kwargs):

~/tfnightly/lib/python3.6/site-packages/tensorflow_probability/python/distributions/distribution.py in _call_sample_n(self, sample_shape, seed, name, **kwargs)
    816       sample_shape, n = self._expand_sample_shape_to_vector(
    817           sample_shape, "sample_shape")
--> 818       samples = self._sample_n(n, seed, **kwargs)
    819       batch_event_shape = tf.shape(input=samples)[1:]
    820       final_shape = tf.concat([sample_shape, batch_event_shape], 0)

~/tfnightly/lib/python3.6/site-packages/tensorflow_probability/python/distributions/mixture.py in _sample_n(self, n, seed)
    374           data=samples_raw_indices,
    375           partitions=cat_samples,
--> 376           num_partitions=self.num_components)
    377 
    378       # Copy the batch indices n times, as we will need to know

~/tfnightly/lib/python3.6/site-packages/tensorflow/python/ops/gen_data_flow_ops.py in dynamic_partition(data, partitions, num_partitions, name)
    718       else:
    719         message = e.message
--> 720       _six.raise_from(_core._status_to_exception(e.code, message), None)
    721   # Add nodes to the TensorFlow graph.
    722   num_partitions = _execute.make_int(num_partitions, "num_partitions")

/usr/lib/python3/dist-packages/six.py in raise_from(value, from_value)

NotFoundError: No registered 'DynamicPartition' OpKernel for GPU devices compatible with node {{node DynamicPartition}}
	 (OpKernel was found, but attributes didn't match) Requested Attributes: T=DT_INT32, num_partitions=2
	.  Registered:  device='CPU'; T in [DT_VARIANT]
  device='CPU'; T in [DT_RESOURCE]
  device='CPU'; T in [DT_STRING]
  device='CPU'; T in [DT_BOOL]
  device='CPU'; T in [DT_COMPLEX128]
  device='CPU'; T in [DT_COMPLEX64]
  device='CPU'; T in [DT_DOUBLE]
  device='CPU'; T in [DT_FLOAT]
  device='CPU'; T in [DT_BFLOAT16]
  device='CPU'; T in [DT_HALF]
  device='CPU'; T in [DT_INT8]
  device='CPU'; T in [DT_UINT8]
  device='CPU'; T in [DT_INT16]
  device='CPU'; T in [DT_UINT16]
  device='CPU'; T in [DT_INT32]
  device='CPU'; T in [DT_INT64]
  device='GPU'; T in [DT_COMPLEX128]
  device='GPU'; T in [DT_COMPLEX64]
  device='GPU'; T in [DT_DOUBLE]
  device='GPU'; T in [DT_FLOAT]
  device='GPU'; T in [DT_HALF]
 [Op:DynamicPartition]

Start in tensorflow_probability/python/distributions/mixture.py at _sample_n, where the traceback reaches DynamicPartition. Reproduce the provided tfd.Mixture sample call on /gpu:0 and investigate the reported missing GPU kernel; done means the sampling example completes successfully on GPU.