Inference Gym: Add 64 bit support to all targets.

This is done via adding a `dtype` argument to the initializer which is then used to cast the inputs (if any) as well as the dtype for de novo arrays (like tf.zeros). This is a departure from the usual TensorFlow Probability style where the dtype is inferred from the inputs because:

- Inference Gym targets often have no inputs
- Inference Gym is not designed to do deferred array materialization like TFP is

We also depart from the typical JAX style of using the maximum precision available because we want to enable testing 32 and 64 bit implementations in the same process.

NOTE: At this time, ground truth remains always numpy arrays and always 64 bit.

Fixes #1993

PiperOrigin-RevId: 736231310

Author: SiegeLordEx

spinoffs/inference_gym/inference_gym/internal/test_util.py

@@ -23,6 +23,7 @@
 import tensorflow.compat.v2 as tf
 import tensorflow_probability as tfp
 
+from tensorflow_probability.python.internal import dtype_util
 from tensorflow_probability.python.internal import test_util
 
 flags.DEFINE_bool('use_tfds', False, 'Whether to run tests that use TFDS.',
@@ -279,6 +280,7 @@ def validate_log_prob_and_transforms(
       self,
       model,
       sample_transformation_shapes,
+      dtype=tf.float32,
       check_ground_truth_mean=False,
       check_ground_truth_mean_standard_error=False,
       check_ground_truth_standard_deviation=False,
@@ -295,6 +297,7 @@ def validate_log_prob_and_transforms(
     Args:
       model: The model to validate.
       sample_transformation_shapes: Shapes of the transformation outputs.
+      dtype: The expected dtype of floating point quantities.
       check_ground_truth_mean: Whether to check the shape of the ground truth
         mean.
       check_ground_truth_mean_standard_error: Whether to check the shape of the
@@ -331,16 +334,25 @@ def _random_element(shape, dtype, default_event_space_bijector, seed):
 
     self.assertAllFinite(log_prob)
     self.assertEqual((batch_size,), log_prob.shape)
+    self.assertEqual(dtype, log_prob.dtype)
+
+    def _assert_dtype_part(part):
+      if dtype_util.is_floating(part):
+        self.assertEqual(dtype, part)
+
+    self.assertAllAssertsNested(_assert_dtype_part, model.dtype)
 
     for name, sample_transformation in model.sample_transformations.items():
-      transformed_points = self.evaluate(sample_transformation(test_points))
+      transformed_points = sample_transformation(test_points)
 
       def _assertions_part(expected_shape, expected_dtype, transformed_part):
         self.assertAllFinite(transformed_part)
         self.assertEqual(
             (batch_size,) + tuple(expected_shape),
             tuple(list(transformed_part.shape)))
         self.assertEqual(expected_dtype, transformed_part.dtype)
+        if dtype_util.is_floating(transformed_part.dtype):
+          self.assertEqual(dtype, transformed_part.dtype)
 
       self.assertAllAssertsNested(
           _assertions_part,

spinoffs/inference_gym/inference_gym/targets/BUILD

@@ -63,6 +63,7 @@ py_library(
     srcs = ["banana.py"],
     deps = [
         ":model",
+        # absl/testing:parameterized dep,
         # numpy dep,
     ],
 )
@@ -73,6 +74,7 @@ py_test(
     srcs = ["banana_test.py"],
     deps = [
         ":banana",
+        # absl/testing:parameterized dep,
         # tensorflow dep,
         # tensorflow_probability/python/internal:test_util dep,
         "//inference_gym/internal:test_util",
@@ -146,6 +148,7 @@ py_test(
     srcs = ["eight_schools_test.py"],
     deps = [
         ":eight_schools",
+        # absl/testing:parameterized dep,
         # tensorflow dep,
         # tensorflow_probability/python/internal:test_util dep,
         "//inference_gym/internal:test_util",
@@ -167,6 +170,7 @@ py_test(
     srcs = ["ill_conditioned_gaussian_test.py"],
     deps = [
         ":ill_conditioned_gaussian",
+        # absl/testing:parameterized dep,
         # tensorflow dep,
         # tensorflow_probability/python/internal:test_util dep,
         "//inference_gym/internal:test_util",
@@ -220,6 +224,7 @@ py_test(
     shard_count = 3,
     deps = [
         ":log_gaussian_cox_process",
+        # absl/testing:parameterized dep,
         # numpy dep,
         # tensorflow dep,
         # tensorflow_probability/python/internal:test_util dep,
@@ -275,6 +280,7 @@ py_test(
     shard_count = 3,
     deps = [
         ":lorenz_system",
+        # absl/testing:parameterized dep,
         # numpy dep,
         # tensorflow dep,
         # tensorflow_probability/python/internal:test_util dep,
@@ -319,6 +325,7 @@ py_test(
     srcs = ["neals_funnel_test.py"],
     deps = [
         ":neals_funnel",
+        # absl/testing:parameterized dep,
         # tensorflow dep,
         # tensorflow_probability/python/internal:test_util dep,
         "//inference_gym/internal:test_util",
@@ -340,6 +347,7 @@ py_test(
     srcs = ["non_identifiable_quartic_test.py"],
     deps = [
         ":non_identifiable_quartic",
+        # absl/testing:parameterized dep,
         # tensorflow dep,
         # tensorflow_probability/python/internal:test_util dep,
         "//inference_gym/internal:test_util",

spinoffs/inference_gym/inference_gym/targets/banana.py

@@ -61,16 +61,18 @@ def __init__(
       self,
       ndims=2,
       curvature=0.03,
+      dtype=tf.float32,
       name='banana',
       pretty_name='Banana',
   ):
-    """Construct the banana model.
+    """Initialize the banana model.
 
     Args:
       ndims: Python integer. Dimensionality of the distribution. Must be at
         least 2.
       curvature: Python float. Controls the strength of the curvature of
         the distribution.
+      dtype: Dtype to use for floating point quantities.
       name: Python `str` name prefixed to Ops created by this class.
       pretty_name: A Python `str`. The pretty name of this model.
 
@@ -87,16 +89,20 @@ def bijector_fn(x):
         batch_shape = ps.shape(x)[:-1]
         shift = tf.concat(
             [
-                tf.zeros(ps.concat([batch_shape, [1]], axis=0)),
+                tf.zeros(ps.concat([batch_shape, [1]], axis=0), dtype),
                 curvature * (tf.square(x[..., :1]) - 100),
-                tf.zeros(ps.concat([batch_shape, [ndims - 2]], axis=0)),
+                tf.zeros(
+                    ps.concat([batch_shape, [ndims - 2]], axis=0), dtype
+                ),
             ],
             axis=-1,
         )
         return tfb.Shift(shift)
 
       mg = tfd.MultivariateNormalDiag(
-          loc=tf.zeros(ndims), scale_diag=[10.] + [1.] * (ndims - 1))
+          loc=tf.zeros(ndims, dtype),
+          scale_diag=[10.0] + [1.0] * (ndims - 1),
+      )
       banana = tfd.TransformedDistribution(
           mg, bijector=tfb.MaskedAutoregressiveFlow(bijector_fn=bijector_fn))
 
@@ -115,6 +121,7 @@ def bijector_fn(x):
                   ground_truth_standard_deviation=np.array(
                       [10.] + [np.sqrt(1. + 2 * curvature**2 * 10.**4)] +
                       [1.] * (ndims - 2)),
+                  dtype=banana.dtype,
               )
       }
 

spinoffs/inference_gym/inference_gym/targets/banana_test.py

@@ -14,6 +14,9 @@
 # ============================================================================
 """Tests for inference_gym.targets.banana."""
 
+from absl.testing import parameterized
+import tensorflow.compat.v2 as tf
+
 from tensorflow_probability.python.internal import test_util as tfp_test_util
 from inference_gym.internal import test_util
 from inference_gym.targets import banana
@@ -22,18 +25,23 @@
 @test_util.multi_backend_test(globals(), 'targets.banana_test')
 class BananaTest(test_util.InferenceGymTestCase):
 
-  def testBasic(self):
+  @parameterized.parameters(tf.float32, tf.float64)
+  def testBasic(self, dtype):
     """Checks that you get finite values given unconstrained samples.
 
     We check `unnormalized_log_prob` as well as the values of the sample
     transformations.
+
+    Args:
+      dtype: Dtype to use for floating point computations.
     """
-    model = banana.Banana(ndims=3)
+    model = banana.Banana(ndims=3, dtype=dtype)
     self.validate_log_prob_and_transforms(
         model,
         sample_transformation_shapes=dict(identity=[3]),
         check_ground_truth_mean=True,
         check_ground_truth_standard_deviation=True,
+        dtype=dtype,
     )
 
   def testMC(self):

spinoffs/inference_gym/inference_gym/targets/brownian_motion.py

@@ -59,24 +59,30 @@ def brownian_motion_prior_fn(num_timesteps,
                            name='x_{}'.format(t))
 
 
-def brownian_motion_unknown_scales_prior_fn(num_timesteps, use_markov_chain):
+def brownian_motion_unknown_scales_prior_fn(
+    num_timesteps, use_markov_chain, dtype
+):
   """Generative process for the Brownian Motion model with unknown scales."""
-  innovation_noise_scale = yield Root(tfd.LogNormal(
-      0., 2., name='innovation_noise_scale'))
-  _ = yield Root(tfd.LogNormal(0., 2., name='observation_noise_scale'))
+  zero = tf.zeros([], dtype)
+  innovation_noise_scale = yield Root(
+      tfd.LogNormal(zero, 2.0, name='innovation_noise_scale')
+  )
+  _ = yield Root(tfd.LogNormal(zero, 2.0, name='observation_noise_scale'))
   if use_markov_chain:
     yield brownian_motion_as_markov_chain(
         num_timesteps=num_timesteps,
-        innovation_noise_scale=innovation_noise_scale)
+        innovation_noise_scale=innovation_noise_scale,
+    )
   else:
     yield from brownian_motion_prior_fn(
-        num_timesteps,
-        innovation_noise_scale=innovation_noise_scale)
+        num_timesteps, innovation_noise_scale=innovation_noise_scale
+    )
 
 
 def brownian_motion_log_likelihood_fn(values,
                                       observed_locs,
                                       use_markov_chain,
+                                      dtype,
                                       observation_noise_scale=None):
   """Likelihood of observed data under the Brownian Motion model."""
   if observation_noise_scale is None:
@@ -86,7 +92,12 @@ def brownian_motion_log_likelihood_fn(values,
     latents = values if use_markov_chain else tf.stack(values, axis=-1)
 
   observation_noise_scale = tf.convert_to_tensor(
-      observation_noise_scale, name='observation_noise_scale')
+      observation_noise_scale, dtype=dtype, name='observation_noise_scale')
+  observed_locs = tf.cast(
+      observed_locs,
+      dtype=dtype,
+      name='observed_locs',
+  )
   is_observed = ~tf.math.is_nan(observed_locs)
   lps = tfd.Normal(
       loc=latents, scale=observation_noise_scale[..., tf.newaxis]).log_prob(
@@ -117,6 +128,7 @@ def __init__(self,
                innovation_noise_scale,
                observation_noise_scale,
                use_markov_chain=False,
+               dtype=tf.float32,
                name='brownian_motion',
                pretty_name='Brownian Motion'):
     """Construct the Brownian Motion model.
@@ -130,11 +142,18 @@ def __init__(self,
         `MarkovChain` distribution in place of separate random variables for
         each time step. The default of `False` is for backwards compatibility;
         setting this to `True` should significantly improve performance.
+      dtype: Dtype to use for floating point quantities.
       name: Python `str` name prefixed to Ops created by this class.
       pretty_name: A Python `str`. The pretty name of this model.
     """
     with tf.name_scope(name):
       num_timesteps = observed_locs.shape[0]
+      innovation_noise_scale = tf.convert_to_tensor(
+          innovation_noise_scale,
+          dtype=dtype,
+          name='innovation_noise_scale',
+      )
+
       if use_markov_chain:
         self._prior_dist = brownian_motion_as_markov_chain(
             num_timesteps=num_timesteps,
@@ -150,7 +169,8 @@ def __init__(self,
           brownian_motion_log_likelihood_fn,
           observation_noise_scale=observation_noise_scale,
           observed_locs=observed_locs,
-          use_markov_chain=use_markov_chain)
+          use_markov_chain=use_markov_chain,
+          dtype=dtype)
 
       def _ext_identity(params):
         return tf.stack(params, axis=-1)
@@ -164,6 +184,7 @@ def _ext_identity_markov_chain(params):
                   fn=(_ext_identity_markov_chain
                       if use_markov_chain else _ext_identity),
                   pretty_name='Identity',
+                  dtype=dtype,
               )
       }
 
@@ -193,12 +214,13 @@ class BrownianMotionMissingMiddleObservations(BrownianMotion):
 
   GROUND_TRUTH_MODULE = brownian_motion_missing_middle_observations
 
-  def __init__(self, use_markov_chain=False):
+  def __init__(self, use_markov_chain=False, dtype=tf.float32):
     dataset = data.brownian_motion_missing_middle_observations()
     super(BrownianMotionMissingMiddleObservations, self).__init__(
         name='brownian_motion_missing_middle_observations',
         pretty_name='Brownian Motion Missing Middle Observations',
         use_markov_chain=use_markov_chain,
+        dtype=dtype,
         **dataset)
 
 
@@ -226,6 +248,7 @@ class BrownianMotionUnknownScales(bayesian_model.BayesianModel):
   def __init__(self,
                observed_locs,
                use_markov_chain=False,
+               dtype=tf.float32,
                name='brownian_motion_unknown_scales',
                pretty_name='Brownian Motion with Unknown Scales'):
     """Construct the Brownian Motion model with unknown scales.
@@ -238,6 +261,7 @@ def __init__(self,
         each time step. The default of `False` is for backwards compatibility;
         setting this to `True` should significantly improve performance.
         Default value: `False`.
+      dtype: Dtype to use for floating point quantities.
       name: Python `str` name prefixed to Ops created by this class.
       pretty_name: A Python `str`. The pretty name of this model.
     """
@@ -247,12 +271,14 @@ def __init__(self,
           functools.partial(
               brownian_motion_unknown_scales_prior_fn,
               use_markov_chain=use_markov_chain,
-              num_timesteps=num_timesteps))
+              num_timesteps=num_timesteps,
+              dtype=dtype))
 
       self._log_likelihood_fn = functools.partial(
           brownian_motion_log_likelihood_fn,
           use_markov_chain=use_markov_chain,
-          observed_locs=observed_locs)
+          observed_locs=observed_locs,
+          dtype=dtype)
 
       def _ext_identity(params):
         return {'innovation_noise_scale': params[0],
@@ -266,9 +292,9 @@ def _ext_identity(params):
               model.Model.SampleTransformation(
                   fn=_ext_identity,
                   pretty_name='Identity',
-                  dtype={'innovation_noise_scale': tf.float32,
-                         'observation_noise_scale': tf.float32,
-                         'locs': tf.float32})
+                  dtype={'innovation_noise_scale': dtype,
+                         'observation_noise_scale': dtype,
+                         'locs': dtype})
       }
 
     event_space_bijector = type(
@@ -300,12 +326,13 @@ class BrownianMotionUnknownScalesMissingMiddleObservations(
   GROUND_TRUTH_MODULE = (
       brownian_motion_unknown_scales_missing_middle_observations)
 
-  def __init__(self, use_markov_chain=False):
+  def __init__(self, use_markov_chain=False, dtype=tf.float32):
     dataset = data.brownian_motion_missing_middle_observations()
     del dataset['innovation_noise_scale']
     del dataset['observation_noise_scale']
     super(BrownianMotionUnknownScalesMissingMiddleObservations, self).__init__(
         name='brownian_motion_unknown_scales_missing_middle_observations',
         pretty_name='Brownian Motion with Unknown Scales',
         use_markov_chain=use_markov_chain,
+        dtype=dtype,
         **dataset)