Ensure that sample_halton_sequence is jittable.

PiperOrigin-RevId: 565432730

Author: srvasude

tensorflow_probability/python/mcmc/sample_halton_sequence_lib.py

@@ -182,17 +182,14 @@ def sample_halton_sequence(dim,
     # 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.
-    if num_results is not None:
-      num_results = tf.convert_to_tensor(num_results)
     if sequence_indices is not None:
       sequence_indices = tf.convert_to_tensor(sequence_indices)
     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)
+    if num_results is None:
+      num_results = ps.reduce_max(indices)
+    radixes = _PRIMES[0:dim][..., np.newaxis]
+    max_sizes_by_axes = _base_expansion_size(num_results, radixes, dtype)
+    max_size = ps.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
@@ -204,14 +201,13 @@ def sample_halton_sequence(dim,
     # 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])
+    exponents_by_axes = tf.tile([tf.range(max_size, dtype=dtype)], [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,
-                                exponents_by_axes,
-                                tf.constant(0, exponents_by_axes.dtype))
-    weights = radixes ** capped_exponents
+    capped_exponents = tf.where(
+        weight_mask, exponents_by_axes, dtype_util.as_numpy_dtype(dtype)(0.))
+    weights = tf.cast(radixes ** capped_exponents, dtype=dtype)
     # 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
@@ -246,22 +242,22 @@ def sample_halton_sequence(dim,
     zero_correction = samplers.uniform([dim, 1],
                                        seed=zero_correction_seed,
                                        dtype=dtype)
-    zero_correction /= radixes ** max_sizes_by_axes
+    zero_correction /= tf.cast(radixes ** max_sizes_by_axes, dtype)
     return base_values + tf.reshape(zero_correction, [-1])
 
 
 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(coeffs)[-1]
-  radixes = tf.reshape(tf.cast(radixes, dtype=tf.int32), shape=[-1])
-  perms = _get_permutations(num_coeffs, radixes, seed=seed)
+  num_coeffs = ps.shape(coeffs)[-1]
+  perms = _get_permutations(num_coeffs, np.squeeze(radixes, axis=-1), seed=seed)
   perms = tf.reshape(perms, shape=[-1])
+  radixes = tf.reshape(tf.cast(radixes, dtype=tf.int32), 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
+  offsets = radix_offsets + ps.range(num_coeffs, dtype=tf.int32) * radix_sum
   permuted_coeffs = tf.gather(perms, coeffs + offsets)
   return tf.cast(permuted_coeffs, dtype=given_dtype)
 
@@ -291,10 +287,11 @@ def _get_permutations(num_results, dims, seed=None):
   seeds = samplers.split_seed(seed, n=ps.size(dims))
 
   def generate_one(dim, seed):
-    return tf.argsort(samplers.uniform([num_results, dim], seed=seed), axis=-1)
+    return tf.argsort(samplers.uniform(
+        [num_results, dim], seed=seed), axis=-1)
 
   return tf.concat([generate_one(dim, seed)
-                    for dim, seed in zip(tf.unstack(dims), tf.unstack(seeds))],
+                    for dim, seed in zip(dims, tf.unstack(seeds))],
                    axis=-1)
 
 
@@ -325,8 +322,13 @@ def _get_indices(num_results, sequence_indices, dtype, name=None):
   """
   with tf.name_scope(name or 'get_indices'):
     if sequence_indices is None:
-      num_results = tf.cast(num_results, dtype=dtype)
-      sequence_indices = tf.range(num_results, dtype=dtype)
+      np_dtype = dtype_util.as_numpy_dtype(dtype)
+      num_results_ = tf.get_static_value(num_results)
+      if num_results_ is not None:
+        sequence_indices = ps.range(np_dtype(num_results_), dtype=dtype)
+      else:
+        num_results = tf.cast(num_results, dtype=dtype)
+        sequence_indices = ps.range(num_results, dtype=dtype)
     else:
       sequence_indices = tf.cast(sequence_indices, dtype)
 
@@ -338,7 +340,7 @@ def _get_indices(num_results, sequence_indices, dtype, name=None):
     return tf.reshape(indices, [-1, 1, 1])
 
 
-def _base_expansion_size(num, bases):
+def _base_expansion_size(num, bases, dtype):
   """Computes the number of terms in the place value expansion.
 
   Let num = a0 + a1 b + a2 b^2 + ... ak b^k be the place value expansion of
@@ -349,16 +351,23 @@ def _base_expansion_size(num, bases):
     $$k = Floor(log_b (num)) + 1  = Floor( log(num) / log(b)) + 1$$
 
   Args:
-    num: Scalar `Tensor` of dtype either `float32` or `float64`. The number to
+    num: Scalar `Tensor` of dtype either `int32` or `int64`. The number to
       compute the base expansion size of.
     bases: `Tensor` of the same dtype as num. The bases to compute the size
       against.
+    dtype: Return `dtype`.
 
   Returns:
-    Tensor of same dtype and shape as `bases` containing the size of num when
+    Tensor of dtype `dtype` and shape as `bases` containing the size of num when
     written in that base.
   """
-  return tf.floor(tf.math.log(num) / tf.math.log(bases)) + 1
+  num_ = tf.get_static_value(num)
+  if num_ is not None:
+    return (np.floor(np.log(num_) / np.log(bases)) + 1).astype(
+        dtype_util.as_numpy_dtype(dtype))
+
+  return tf.floor(
+      tf.math.log(tf.cast(num, dtype)) / tf.math.log(tf.cast(bases, dtype))) + 1
 
 
 def _primes_less_than(n):

tensorflow_probability/python/mcmc/sample_halton_sequence_test.py

@@ -77,6 +77,23 @@ def test_dtypes_works_correctly(self):
     self.assertEqual(self.evaluate(sample_float32).dtype, np.float32)
     self.assertEqual(self.evaluate(sample_float64).dtype, np.float64)
 
+  @test_util.disable_test_for_backend(
+      disable_numpy=True, reason="Numpy has no notion of jit compilation.")
+  def test_jit_works_correctly(self):
+    @tf.function(jit_compile=True)
+    def sample_float32():
+      return sample_halton_sequence_lib.sample_halton_sequence(
+          5, num_results=10, dtype=tf.float32, seed=test_util.test_seed())
+    samples = sample_float32()
+    self.assertEqual(samples.shape, [10, 5])
+
+    @tf.function(jit_compile=True)
+    def sample_float64():
+      return sample_halton_sequence_lib.sample_halton_sequence(
+          5, num_results=10, dtype=tf.float64, seed=test_util.test_seed())
+    samples = sample_float64()
+    self.assertEqual(samples.shape, [10, 5])
+
   def test_normal_integral_mean_and_var_correctly_estimated(self):
     n = 1000
     # This test is almost identical to the similarly named test in