tfp.substrates.jax.mcmc.sample_halton_sequence

Returns a sample from the dim dimensional Halton sequence.

Computes the members of the low discrepancy Halton sequence in dimension dim. The dim-dimensional sequence takes values in the unit hypercube in dim dimensions. Currently, only dimensions up to 10000 are supported. The prime base for the k-th axes is the k-th prime starting from 2. For example, if dim = 3, then the bases will be [2, 3, 5] respectively and the first element of the non-randomized sequence will be: [0.5, 0.333, 0.2]. For a more complete description of the Halton sequences see here. For low discrepancy sequences and their applications see here.

If randomized is true, this function produces a scrambled version of the Halton sequence introduced by [Owen (2017)][1]. For the advantages of randomization of low discrepancy sequences see here.

The number of samples produced is controlled by the num_results and sequence_indices parameters. The user must supply either num_results or sequence_indices but not both. The former is the number of samples to produce starting from the first element. If sequence_indices is given instead, the specified elements of the sequence are generated. For example, sequence_indices=tf.range(10) is equivalent to specifying n=10.

Examples

from tensorflow_probability.python.internal.backend import jax as tf
import tensorflow_probability as tfp; tfp = tfp.substrates.jax

# Produce the first 1000 members of the Halton sequence in 3 dimensions.
num_results = 1000
dim = 3
sample = tfp.mcmc.sample_halton_sequence(
  dim,
  num_results=num_results,
  seed=127)

# Evaluate the integral of x_1 * x_2^2 * x_3^3  over the three dimensional
# hypercube.
powers = tf.range(1.0, limit=dim + 1)
integral = tf.reduce_mean(tf.reduce_prod(sample ** powers, axis=-1))
true_value = 1.0 / tf.reduce_prod(powers + 1.0)
with tf.Session() as session:
  values = session.run((integral, true_value))

# Produces a relative absolute error of 1.7%.
print ("Estimated: %f, True Value: %f" % values)

# Now skip the first 1000 samples and recompute the integral with the next
# thousand samples. The sequence_indices argument can be used to do this.


sequence_indices = tf.range(start=1000, limit=1000 + num_results,
                            dtype=tf.int32)
sample_leaped = tfp.mcmc.sample_halton_sequence(
    dim,
    sequence_indices=sequence_indices,
    seed=111217)

integral_leaped = tf.reduce_mean(tf.reduce_prod(sample_leaped ** powers,
                                                axis=-1))
with tf.Session() as session:
  values = session.run((integral_leaped, true_value))
# Now produces a relative absolute error of 0.05%.
print ("Leaped Estimated: %f, True Value: %f" % values)

dim Positive Python int representing each sample's event_size. Must not be greater than 10000.
num_results (Optional) Positive scalar Tensor of dtype int32. The number of samples to generate. Either this parameter or sequence_indices must be specified but not both. If this parameter is None, then the behaviour is determined by the sequence_indices. Default value: None.
sequence_indices (Optional) Tensor of dtype int32 and rank 1. The elements of the sequence to compute specified by their position in the sequence. The entries index into the Halton sequence starting with 0 and hence, must be whole numbers. For example, sequence_indices=[0, 5, 6] will produce the first, sixth and seventh elements of the sequence. If this parameter is None, then the num_results parameter must be specified which gives the number of desired samples starting from the first sample. Default value: None.
dtype (Optional) The dtype of the sample. One of: float16, float32 or float64. Default value: tf.float32.
randomized (Optional) bool indicating whether to produce a randomized Halton sequence. If True, applies the randomization described in [Owen (2017)][1]. Default value: True.
seed PRNG seed; see tfp.random.sanitize_seed for details. Only 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".

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.

ValueError if both sequence_indices and num_results were specified or if dimension dim is less than 1 or greater than 10000.

References

[1]: Art B. Owen. A randomized Halton algorithm in R. arXiv preprint arXiv:1706.02808, 2017. https://arxiv.org/abs/1706.02808