tensorflow_probability/python/mcmc/sample_halton_sequence.py
def _randomize(coeffs, radixes, seed=None):
  """Applies the Owen (2017) randomization to the coefficients."""
  given_dtype = coeffs.dtype
  coeffs = tf.cast(coeffs, dtype=tf.int32)
  num_coeffs = tf.shape(input=coeffs)[-1]
  radixes = tf.reshape(tf.cast(radixes, dtype=tf.int32), shape=[-1])
  stream = distributions.SeedStream(seed, salt='MCMCSampleHaltonSequence2')
  perms = _get_permutations(num_coeffs, radixes, seed=stream())
  perms = tf.reshape(perms, shape=[-1])
  radix_sum = tf.reduce_sum(input_tensor=radixes)
  radix_offsets = tf.reshape(tf.cumsum(radixes, exclusive=True),
                             shape=[-1, 1])
  offsets = radix_offsets + tf.range(num_coeffs) * radix_sum
  permuted_coeffs = tf.gather(perms, coeffs + offsets)
  return tf.cast(permuted_coeffs, dtype=given_dtype)
def _randomize(coeffs, radixes, seed=None):
  """Applies the Owen (2017) randomization to the coefficients."""
  given_dtype = coeffs.dtype
  coeffs = tf.cast(coeffs, dtype=tf.int32)
  num_coeffs = tf.shape(input=coeffs)[-1]
  radixes = tf.reshape(tf.cast(radixes, dtype=tf.int32), shape=[-1])
  stream = distributions.SeedStream(seed, salt='MCMCSampleHaltonSequence2')
  perms = _get_permutations(num_coeffs, radixes, seed=stream())
  perms = tf.reshape(perms, shape=[-1])
  radix_sum = tf.reduce_sum(input_tensor=radixes)
  radix_offsets = tf.reshape(tf.cumsum(radixes, exclusive=True),
                             shape=[-1, 1])
  offsets = radix_offsets + tf.range(num_coeffs) * radix_sum
  permuted_coeffs = tf.gather(perms, coeffs + offsets)
  return tf.cast(permuted_coeffs, dtype=given_dtype)
Applies the Owen (2017) randomization to the coefficients.