def default_exchange_proposed_fn_(num_replica, seed=None):
    """Default function for `exchange_proposed_fn` of `kernel`."""
    num_replica = tf.to_int32(num_replica)

    seed = distributions_util.gen_new_seed(seed, 'default_exchange_proposed_fn')
    random_uniform = tf.random_uniform([], seed=seed)
    accept_proposed_exchange = random_uniform < probs

    seed = distributions_util.gen_new_seed(seed, 'default_exchange_proposed_fn')
    zero_start = tf.random_uniform([], seed=seed) > 0.5
    if num_replica % 2 == 0:
      exchange_proposed = tf.where(
          zero_start, tf.range(num_replica),
          tf.sparse_to_dense(tf.range(num_replica - 2), (num_replica,),
                             tf.range(1, num_replica - 1)))
      exchange_proposed_n = tf.where(zero_start, num_replica // 2,
                                     num_replica // 2 - 1)
    else:
      exchange_proposed = tf.where(
          zero_start, tf.range(num_replica - 1), tf.range(1, num_replica))
      exchange_proposed_n = num_replica // 2

    exchange_proposed = tf.reshape(exchange_proposed, (num_replica // 2, 2))
    exchange_proposed = tf.where(accept_proposed_exchange, exchange_proposed,
                                 tf.zeros_like(exchange_proposed))
    exchange_proposed_n = tf.where(accept_proposed_exchange,
                                   exchange_proposed_n,
                                   tf.zeros_like(exchange_proposed_n))
    return exchange_proposed, exchange_proposed_n

 def body(i, next_replica_idx):
   """`tf.while_loop` body."""
   ratio = (
       sampled_replica_ratios[next_replica_idx[exchange_proposed[i, 0]]]
       - sampled_replica_ratios[next_replica_idx[exchange_proposed[i, 1]]])
   ratio *= (
       self.inverse_temperatures[exchange_proposed[i, 1]]
       - self.inverse_temperatures[exchange_proposed[i, 0]])
   self._seed_stream = distributions_util.gen_new_seed(
       self._seed_stream, salt='replica_exchange_one_step')
   log_uniform = tf.log(tf.random_uniform(
       shape=tf.shape(ratio),
       dtype=ratio.dtype.base_dtype,
       seed=self._seed_stream))
   exchange = log_uniform < ratio
   exchange_op = tf.sparse_to_dense(
       [exchange_proposed[i, 0], exchange_proposed[i, 1]],
       [self.num_replica],
       [next_replica_idx[exchange_proposed[i, 1]] -
        next_replica_idx[exchange_proposed[i, 0]],
        next_replica_idx[exchange_proposed[i, 0]] -
        next_replica_idx[exchange_proposed[i, 1]]])
   next_replica_idx = tf.cond(exchange,
                              lambda: next_replica_idx + exchange_op,
                              lambda: next_replica_idx)
   return [i + 1, next_replica_idx]

 def _sample_n(self, n, seed=None):
   # Get ids as a [n, batch_size]-shaped matrix, unless batch_shape=[] then get
   # ids as a [n]-shaped vector.
   batch_size = (np.prod(self.batch_shape.as_list(), dtype=np.int32)
                 if self.batch_shape.is_fully_defined()
                 else math_ops.reduce_prod(self.batch_shape_tensor()))
   ids = self._mixture_distribution.sample(
       sample_shape=concat_vectors(
           [n],
           distribution_util.pick_vector(
               self.is_scalar_batch(),
               np.int32([]),
               [batch_size])),
       seed=distribution_util.gen_new_seed(
           seed, "poisson_lognormal_quadrature_compound"))
   # Stride `quadrature_degree` for `batch_size` number of times.
   offset = math_ops.range(start=0,
                           limit=batch_size * len(self.quadrature_probs),
                           delta=len(self.quadrature_probs),
                           dtype=ids.dtype)
   ids += offset
   rate = array_ops.gather(
       array_ops.reshape(self.distribution.rate, shape=[-1]), ids)
   rate = array_ops.reshape(
       rate, shape=concat_vectors([n], self.batch_shape_tensor()))
   return random_ops.random_poisson(
       lam=rate, shape=[], dtype=self.dtype, seed=seed)

  def one_step(self, current_state, previous_kernel_results):
    with tf.name_scope(
        name=mcmc_util.make_name(self.name, 'rwm', 'one_step'),
        values=[self.seed,
                current_state,
                previous_kernel_results.target_log_prob]):
      with tf.name_scope('initialize'):
        current_state_parts = (list(current_state)
                               if mcmc_util.is_list_like(current_state)
                               else [current_state])
        current_state_parts = [tf.convert_to_tensor(s, name='current_state')
                               for s in current_state_parts]

      self._seed_stream = distributions_util.gen_new_seed(
          self._seed_stream, salt='rwm_kernel_proposal')
      new_state_fn = self.new_state_fn
      next_state_parts = new_state_fn(current_state_parts, self._seed_stream)
      # Compute `target_log_prob` so its available to MetropolisHastings.
      next_target_log_prob = self.target_log_prob_fn(*next_state_parts)

      def maybe_flatten(x):
        return x if mcmc_util.is_list_like(current_state) else x[0]

      return [
          maybe_flatten(next_state_parts),
          UncalibratedRandomWalkResults(
              log_acceptance_correction=tf.zeros(
                  shape=tf.shape(next_target_log_prob),
                  dtype=next_target_log_prob.dtype.base_dtype),
              target_log_prob=next_target_log_prob,
          ),
      ]

  def _apply_variational_kernel(self, inputs):
    if (not isinstance(self.kernel_posterior, tfd.Independent) or
        not isinstance(self.kernel_posterior.distribution, tfd.Normal)):
      raise TypeError(
          '`DenseFlipout` requires '
          '`kernel_posterior_fn` produce an instance of '
          '`tf.distributions.Independent(tf.distributions.Normal)` '
          '(saw: \"{}\").'.format(self.kernel_posterior.name))
    self.kernel_posterior_affine = tfd.Normal(
        loc=tf.zeros_like(self.kernel_posterior.distribution.loc),
        scale=self.kernel_posterior.distribution.scale)
    self.kernel_posterior_affine_tensor = (
        self.kernel_posterior_tensor_fn(self.kernel_posterior_affine))
    self.kernel_posterior_tensor = None

    input_shape = tf.shape(inputs)
    batch_shape = input_shape[:-1]

    sign_input = random_rademacher(
        input_shape,
        dtype=inputs.dtype,
        seed=self.seed)
    sign_output = random_rademacher(
        tf.concat([batch_shape,
                   tf.expand_dims(self.units, 0)], 0),
        dtype=inputs.dtype,
        seed=distribution_util.gen_new_seed(
            self.seed, salt='dense_flipout'))
    perturbed_inputs = self._matmul(
        inputs * sign_input, self.kernel_posterior_affine_tensor) * sign_output

    outputs = self._matmul(inputs, self.kernel_posterior.distribution.loc)
    outputs += perturbed_inputs
    return outputs

 def generate_one(d):
   seed[0] = distributions_util.gen_new_seed(
       seed[0], salt='mcmc_sample_halton_sequence_4')
   fn = lambda _: tf.random_shuffle(tf.range(d), seed=seed[0])
   return tf.map_fn(
       fn,
       sample_range,
       parallel_iterations=1 if seed[0] is not None else 10)

 def _fn(state_parts, seed):
   next_state_parts = []
   for state in state_parts:
     # Mutate seed with each use.
     seed = distributions_util.gen_new_seed(
         seed, salt='random_walk_cauchy_increment')
     next_state_parts.append(state + cauchy.sample(
         sample_shape=state.shape, seed=seed))
   return next_state_parts

 def _sample_n(self, n, seed=None):
   if seed is None:
     seed = distribution_util.gen_new_seed(
         seed=np.random.randint(2**32 - 1),
         salt="autoregressive")
   samples = self.distribution0.sample(n, seed=seed)
   for _ in range(self._num_steps):
     samples = self.distribution_fn(samples).sample(seed=seed)
   return samples

  def _sample_n(self, n, seed):
    batch_shape = self.batch_shape_tensor()
    event_shape = self.event_shape_tensor()
    batch_ndims = array_ops.shape(batch_shape)[0]

    ndims = batch_ndims + 3  # sample_ndims=1, event_ndims=2
    shape = array_ops.concat([[n], batch_shape, event_shape], 0)

    # Complexity: O(nbk**2)
    x = random_ops.random_normal(shape=shape,
                                 mean=0.,
                                 stddev=1.,
                                 dtype=self.dtype,
                                 seed=seed)

    # Complexity: O(nbk)
    # This parametrization is equivalent to Chi2, i.e.,
    # ChiSquared(k) == Gamma(alpha=k/2, beta=1/2)
    g = random_ops.random_gamma(shape=[n],
                                alpha=self._multi_gamma_sequence(
                                    0.5 * self.df, self.dimension),
                                beta=0.5,
                                dtype=self.dtype,
                                seed=distribution_util.gen_new_seed(
                                    seed, "wishart"))

    # Complexity: O(nbk**2)
    x = array_ops.matrix_band_part(x, -1, 0)  # Tri-lower.

    # Complexity: O(nbk)
    x = array_ops.matrix_set_diag(x, math_ops.sqrt(g))

    # Make batch-op ready.
    # Complexity: O(nbk**2)
    perm = array_ops.concat([math_ops.range(1, ndims), [0]], 0)
    x = array_ops.transpose(x, perm)
    shape = array_ops.concat([batch_shape, [event_shape[0]], [-1]], 0)
    x = array_ops.reshape(x, shape)

    # Complexity: O(nbM) where M is the complexity of the operator solving a
    # vector system. E.g., for OperatorPDDiag, each matmul is O(k**2), so
    # this complexity is O(nbk**2). For OperatorPDCholesky, each matmul is
    # O(k^3) so this step has complexity O(nbk^3).
    x = self.scale_operator_pd.sqrt_matmul(x)

    # Undo make batch-op ready.
    # Complexity: O(nbk**2)
    shape = array_ops.concat([batch_shape, event_shape, [n]], 0)
    x = array_ops.reshape(x, shape)
    perm = array_ops.concat([[ndims - 1], math_ops.range(0, ndims - 1)], 0)
    x = array_ops.transpose(x, perm)

    if not self.cholesky_input_output_matrices:
      # Complexity: O(nbk^3)
      x = math_ops.matmul(x, x, adjoint_b=True)

    return x

 def _sample_n(self, n, seed=None):
   # Here we use the fact that if:
   # lam ~ Gamma(concentration=total_count, rate=(1-probs)/probs)
   # then X ~ Poisson(lam) is Negative Binomially distributed.
   rate = random_ops.random_gamma(
       shape=[n],
       alpha=self.total_count,
       beta=math_ops.exp(-self.logits),
       dtype=self.dtype,
       seed=seed)
   return random_ops.random_poisson(
       rate,
       shape=[],
       dtype=self.dtype,
       seed=distribution_util.gen_new_seed(seed, "negative_binom"))

def _randomize(coeffs, radixes, seed=None):
  """Applies the Owen (2017) randomization to the coefficients."""
  given_dtype = coeffs.dtype
  coeffs = tf.to_int32(coeffs)
  num_coeffs = tf.shape(coeffs)[-1]
  radixes = tf.reshape(tf.to_int32(radixes), shape=[-1])
  seed = distributions_util.gen_new_seed(
      seed, salt='mcmc_sample_halton_sequence_3')
  perms = _get_permutations(num_coeffs, radixes, seed=seed)
  perms = tf.reshape(perms, shape=[-1])
  radix_sum = tf.reduce_sum(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 __init__(self, target_log_prob_fn, inverse_temperatures,
               make_kernel_fn,
               exchange_proposed_fn=default_exchange_proposed_fn(1.),
               seed=None, name=None, **kwargs):
    """Instantiates this object.

    Args:
      target_log_prob_fn: Python callable which takes an argument like
        `current_state` (or `*current_state` if it's a list) and returns its
        (possibly unnormalized) log-density under the target distribution.
      inverse_temperatures: sequence of inverse temperatures to perform
        samplings with each replica. Must have statically known `rank` and
        statically known leading shape, i.e.,
        `inverse_temperatures.shape[0].value is not None`
      make_kernel_fn: Python callable which takes target_log_prob_fn and seed
        args and returns a TransitionKernel instance.
      exchange_proposed_fn: Python callable which take a number of replicas, and
        return combinations of replicas for exchange and a number of
        combinations.
      seed: Python integer to seed the random number generator.
        Default value: `None` (i.e., no seed).
      name: Python `str` name prefixed to Ops created by this function.
        Default value: `None` (i.e., "remc_kernel").
      **kwargs: Arguments for `make_kernel_fn`.

    Raises:
      ValueError: if `inverse_temperatures` doesn't have statically known rank
        and statically known leading shape
    """
    if inverse_temperatures.shape.ndims is None or \
       inverse_temperatures.shape[0].value is None:
      raise ValueError('"inverse_temperatures" must have statically known rank '
                       'and statically known leading shape')
    self._seed_stream = seed  # This will be mutated with use.
    self._parameters = dict(target_log_prob_fn=target_log_prob_fn,
                            inverse_temperatures=inverse_temperatures,
                            num_replica=inverse_temperatures.shape[0],
                            exchange_proposed_fn=exchange_proposed_fn,
                            seed=seed, name=name)
    self.replica_kernels = []
    for i in range(self.num_replica):
      self._seed_stream = distributions_util.gen_new_seed(
          self._seed_stream, salt='replica_kernels')
      self.replica_kernels.append(make_kernel_fn(
          target_log_prob_fn=_replica_log_prob_fn(
              inverse_temperatures[i], target_log_prob_fn),
          seed=self._seed_stream))

 def _sample_n(self, n, seed=None):
   n_draws = math_ops.cast(self.total_count, dtype=dtypes.int32)
   k = self.event_shape_tensor()[0]
   unnormalized_logits = array_ops.reshape(
       math_ops.log(random_ops.random_gamma(
           shape=[n],
           alpha=self.concentration,
           dtype=self.dtype,
           seed=seed)),
       shape=[-1, k])
   draws = random_ops.multinomial(
       logits=unnormalized_logits,
       num_samples=n_draws,
       seed=distribution_util.gen_new_seed(seed, salt="dirichlet_multinomial"))
   x = math_ops.reduce_sum(array_ops.one_hot(draws, depth=k), -2)
   final_shape = array_ops.concat([[n], self.batch_shape_tensor(), [k]], 0)
   return array_ops.reshape(x, final_shape)

 def _sample_n(self, n, seed=None):
   expanded_concentration1 = array_ops.ones_like(
       self.total_concentration, dtype=self.dtype) * self.concentration1
   expanded_concentration0 = array_ops.ones_like(
       self.total_concentration, dtype=self.dtype) * self.concentration0
   gamma1_sample = random_ops.random_gamma(
       shape=[n],
       alpha=expanded_concentration1,
       dtype=self.dtype,
       seed=seed)
   gamma2_sample = random_ops.random_gamma(
       shape=[n],
       alpha=expanded_concentration0,
       dtype=self.dtype,
       seed=distribution_util.gen_new_seed(seed, "beta"))
   beta_sample = gamma1_sample / (gamma1_sample + gamma2_sample)
   return beta_sample

 def _sample_n(self, n, seed=None):
   # The sampling method comes from the fact that if:
   #   X ~ Normal(0, 1)
   #   Z ~ Chi2(df)
   #   Y = X / sqrt(Z / df)
   # then:
   #   Y ~ StudentT(df).
   shape = array_ops.concat([[n], self.batch_shape_tensor()], 0)
   normal_sample = random_ops.random_normal(shape, dtype=self.dtype, seed=seed)
   df = self.df * array_ops.ones(self.batch_shape_tensor(), dtype=self.dtype)
   gamma_sample = random_ops.random_gamma(
       [n],
       0.5 * df,
       beta=0.5,
       dtype=self.dtype,
       seed=distribution_util.gen_new_seed(seed, salt="student_t"))
   samples = normal_sample * math_ops.rsqrt(gamma_sample / df)
   return samples * self.scale + self.loc  # Abs(scale) not wanted.

  def _fn(state_parts, seed):
    """Adds a normal perturbation to the input state.

    Args:
      state_parts: A list of `Tensor`s of any shape and real dtype representing
        the state parts of the `current_state` of the Markov chain.
      seed: `int` or None. The random seed for this `Op`. If `None`, no seed is
        applied.
        Default value: `None`.

    Returns:
      perturbed_state_parts: A Python `list` of The `Tensor`s. Has the same
        shape and type as the `state_parts`.

    Raises:
      ValueError: if `scale` does not broadcast with `state_parts`.
    """
    with tf.name_scope(name, 'random_walk_normal_fn',
                       values=[state_parts, scale, seed]):
      scales = scale if mcmc_util.is_list_like(scale) else [scale]
      if len(scales) == 1:
        scales *= len(state_parts)
      if len(state_parts) != len(scales):
        raise ValueError('`scale` must broadcast with `state_parts`.')
      next_state_parts = []
      for scale_part, state_part in zip(scales, state_parts):
        # Mutate seed with each use.
        seed = distributions_util.gen_new_seed(
            seed, salt='random_walk_normal_fn')
        next_state_parts.append(tf.random_normal(
            mean=state_part,
            stddev=scale_part,
            shape=tf.shape(state_part),
            dtype=state_part.dtype.base_dtype,
            seed=seed))
      return next_state_parts

  def testDenseFlipout(self):
    batch_size, in_size, out_size = 2, 3, 4
    with self.test_session() as sess:
      (kernel_posterior, kernel_prior, kernel_divergence,
       bias_posterior, bias_prior, bias_divergence, layer, inputs,
       outputs, kl_penalty) = self._testDenseSetUp(
           prob_layers_lib.DenseFlipout,
           batch_size, in_size, out_size, seed=44)

      expected_kernel_posterior_affine = normal_lib.Normal(
          loc=array_ops.zeros_like(kernel_posterior.result_loc),
          scale=kernel_posterior.result_scale)
      expected_kernel_posterior_affine_tensor = (
          expected_kernel_posterior_affine.sample(seed=42))

      sign_input = random_ops.random_uniform(
          [batch_size, in_size],
          minval=0,
          maxval=2,
          dtype=dtypes.int32,
          seed=layer.seed)
      sign_input = math_ops.cast(2 * sign_input - 1, inputs.dtype)
      sign_output = random_ops.random_uniform(
          [batch_size, out_size],
          minval=0,
          maxval=2,
          dtype=dtypes.int32,
          seed=distribution_util.gen_new_seed(
              layer.seed, salt="dense_flipout"))
      sign_output = math_ops.cast(2 * sign_output - 1, inputs.dtype)
      perturbed_inputs = math_ops.matmul(
          inputs * sign_input, expected_kernel_posterior_affine_tensor)
      perturbed_inputs *= sign_output

      expected_outputs = math_ops.matmul(inputs, kernel_posterior.result_loc)
      expected_outputs += perturbed_inputs
      expected_outputs += bias_posterior.result_sample

      [
          expected_outputs_, actual_outputs_,
          expected_kernel_divergence_, actual_kernel_divergence_,
          expected_bias_, actual_bias_,
          expected_bias_divergence_, actual_bias_divergence_,
      ] = sess.run([
          expected_outputs, outputs,
          kernel_divergence.result, kl_penalty[0],
          bias_posterior.result_sample, layer.bias_posterior_tensor,
          bias_divergence.result, kl_penalty[1],
      ])

      self.assertAllClose(
          expected_bias_, actual_bias_,
          rtol=1e-6, atol=0.)
      self.assertAllClose(
          expected_outputs_, actual_outputs_,
          rtol=1e-6, atol=0.)
      self.assertAllClose(
          expected_kernel_divergence_, actual_kernel_divergence_,
          rtol=1e-6, atol=0.)
      self.assertAllClose(
          expected_bias_divergence_, actual_bias_divergence_,
          rtol=1e-6, atol=0.)

      self.assertAllEqual(
          [[kernel_posterior.distribution, kernel_prior.distribution, None]],
          kernel_divergence.args)

      self.assertAllEqual(
          [[bias_posterior.distribution,
            bias_prior.distribution,
            bias_posterior.result_sample]],
          bias_divergence.args)

 def init_momentum(s):
   return random_ops.random_normal(
       shape=array_ops.shape(s),
       dtype=s.dtype.base_dtype,
       seed=distributions_util.gen_new_seed(
           seed, salt="hmc_kernel_momentums"))

#.........这里部分代码省略.........
      Default value: `None` (i.e., compute as needed).
    name: Python `str` name prefixed to Ops created by this function.
      Default value: `None` (i.e., "hmc_kernel").

  Returns:
    accepted_state: Tensor or Python list of `Tensor`s representing the state(s)
      of the Markov chain(s) at each result step. Has same shape as
      `current_state`.
    kernel_results: `collections.namedtuple` of internal calculations used to
      advance the chain.

  Raises:
    ValueError: if there isn't one `step_size` or a list with same length as
      `current_state`.
  """
  with ops.name_scope(
      name, "hmc_kernel",
      [current_state, step_size, num_leapfrog_steps, seed,
       current_target_log_prob, current_grads_target_log_prob]):
    with ops.name_scope("initialize"):
      [current_state_parts, step_sizes, current_target_log_prob,
       current_grads_target_log_prob] = _prepare_args(
           target_log_prob_fn, current_state, step_size,
           current_target_log_prob, current_grads_target_log_prob,
           maybe_expand=True)
      independent_chain_ndims = distributions_util.prefer_static_rank(
          current_target_log_prob)
      current_momentums = []
      for s in current_state_parts:
        current_momentums.append(random_ops.random_normal(
            shape=array_ops.shape(s),
            dtype=s.dtype.base_dtype,
            seed=seed))
        seed = distributions_util.gen_new_seed(
            seed, salt="hmc_kernel_momentums")

      num_leapfrog_steps = ops.convert_to_tensor(
          num_leapfrog_steps,
          dtype=dtypes.int32,
          name="num_leapfrog_steps")
    [
        proposed_momentums,
        proposed_state_parts,
        proposed_target_log_prob,
        proposed_grads_target_log_prob,
    ] = _leapfrog_integrator(current_momentums,
                             target_log_prob_fn,
                             current_state_parts,
                             step_sizes,
                             num_leapfrog_steps,
                             current_target_log_prob,
                             current_grads_target_log_prob)

    energy_change = _compute_energy_change(current_target_log_prob,
                                           current_momentums,
                                           proposed_target_log_prob,
                                           proposed_momentums,
                                           independent_chain_ndims)

    # u < exp(min(-energy, 0)),  where u~Uniform[0,1)
    # ==> -log(u) >= max(e, 0)
    # ==> -log(u) >= e
    # (Perhaps surprisingly, we don't have a better way to obtain a random
    # uniform from positive reals, i.e., `tf.random_uniform(minval=0,
    # maxval=np.inf)` won't work.)
    random_uniform = random_ops.random_uniform(

#.........这里部分代码省略.........
      used if `randomized` is True. If not supplied and `randomized` is True,
      no seed is set.
      Default value: `None`.
    name:  (Optional) Python `str` describing ops managed by this function. If
      not supplied the name of this function is used.
      Default value: "sample_halton_sequence".

  Returns:
    halton_elements: Elements of the Halton sequence. `Tensor` of supplied dtype
      and `shape` `[num_results, dim]` if `num_results` was specified or shape
      `[s, dim]` where s is the size of `sequence_indices` if `sequence_indices`
      were specified.

  Raises:
    ValueError: if both `sequence_indices` and `num_results` were specified or
      if dimension `dim` is less than 1 or greater than 1000.

  #### References

  [1]: Art B. Owen. A randomized Halton algorithm in R. _arXiv preprint
       arXiv:1706.02808_, 2017. https://arxiv.org/abs/1706.02808
  """
  if dim < 1 or dim > _MAX_DIMENSION:
    raise ValueError(
        'Dimension must be between 1 and {}. Supplied {}'.format(_MAX_DIMENSION,
                                                                 dim))
  if (num_results is None) == (sequence_indices is None):
    raise ValueError('Either `num_results` or `sequence_indices` must be'
                     ' specified but not both.')

  if not dtype.is_floating:
    raise ValueError('dtype must be of `float`-type')

  with tf.name_scope(name, 'sample', values=[sequence_indices]):
    # Here and in the following, the shape layout is as follows:
    # [sample dimension, event dimension, coefficient dimension].
    # The coefficient dimension is an intermediate axes which will hold the
    # weights of the starting integer when expressed in the (prime) base for
    # an event dimension.
    indices = _get_indices(num_results, sequence_indices, dtype)
    radixes = tf.constant(_PRIMES[0:dim], dtype=dtype, shape=[dim, 1])

    max_sizes_by_axes = _base_expansion_size(tf.reduce_max(indices),
                                             radixes)

    max_size = tf.reduce_max(max_sizes_by_axes)

    # The powers of the radixes that we will need. Note that there is a bit
    # of an excess here. Suppose we need the place value coefficients of 7
    # in base 2 and 3. For 2, we will have 3 digits but we only need 2 digits
    # for base 3. However, we can only create rectangular tensors so we
    # store both expansions in a [2, 3] tensor. This leads to the problem that
    # we might end up attempting to raise large numbers to large powers. For
    # example, base 2 expansion of 1024 has 10 digits. If we were in 10
    # dimensions, then the 10th prime (29) we will end up computing 29^10 even
    # though we don't need it. We avoid this by setting the exponents for each
    # axes to 0 beyond the maximum value needed for that dimension.
    exponents_by_axes = tf.tile([tf.range(max_size)], [dim, 1])

    # The mask is true for those coefficients that are irrelevant.
    weight_mask = exponents_by_axes >= max_sizes_by_axes
    capped_exponents = tf.where(
        weight_mask,
        tf.zeros_like(exponents_by_axes),
        exponents_by_axes)
    weights = radixes ** capped_exponents
    # The following computes the base b expansion of the indices. Suppose,
    # x = a0 + a1*b + a2*b^2 + ... Then, performing a floor div of x with
    # the vector (1, b, b^2, b^3, ...) will produce
    # (a0 + s1 * b, a1 + s2 * b, ...) where s_i are coefficients we don't care
    # about. Noting that all a_i < b by definition of place value expansion,
    # we see that taking the elements mod b of the above vector produces the
    # place value expansion coefficients.
    coeffs = tf.floor_div(indices, weights)
    coeffs *= 1. - tf.cast(weight_mask, dtype)
    coeffs %= radixes
    if not randomized:
      coeffs /= radixes
      return tf.reduce_sum(coeffs / weights, axis=-1)
    seed = distributions_util.gen_new_seed(
        seed, salt='mcmc_sample_halton_sequence_1')
    coeffs = _randomize(coeffs, radixes, seed=seed)
    # Remove the contribution from randomizing the trailing zero for the
    # axes where max_size_by_axes < max_size. This will be accounted
    # for separately below (using zero_correction).
    coeffs *= 1. - tf.cast(weight_mask, dtype)
    coeffs /= radixes
    base_values = tf.reduce_sum(coeffs / weights, axis=-1)

    # The randomization used in Owen (2017) does not leave 0 invariant. While
    # we have accounted for the randomization of the first `max_size_by_axes`
    # coefficients, we still need to correct for the trailing zeros. Luckily,
    # this is equivalent to adding a uniform random value scaled so the first
    # `max_size_by_axes` coefficients are zero. The following statements perform
    # this correction.
    seed = distributions_util.gen_new_seed(
        seed, salt='mcmc_sample_halton_sequence_2')
    zero_correction = tf.random_uniform([dim, 1], seed=seed, dtype=dtype)
    zero_correction /= radixes ** max_sizes_by_axes
    return base_values + tf.reshape(zero_correction, [-1])