class_docstr: A TestCase mixin for common tests on inference gym targets.
class_docstr_tok: ["A","TestCase","mixin","for","common","tests","on","inference","gym","targets","."]
docstr: Validate that the model's log probability and sample transformations run.
  
  This checks that unconstrained values passed through the event space
  bijectors into `unnormalized_log_prob` and sample transformations yield
  finite values. This also verifies that the transformed values have the
  expected shape.
  
  Args:
    model: The model to validate.
    sample_transformation_shapes: Shapes of the transformation outputs.
    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
      ground truth standard error.
    check_ground_truth_standard_deviation: Whether to check the shape of the
      ground truth standard deviation.
    check_ground_truth_standard_deviation_standard_error: Whether to check the
      shape of the ground truth standard deviation standard error.
    seed: Optional seed to use. By default, `test_util.test_seed()` is used.
docstr_tok: ["Validate","that","the","model","'s","log","probability","and","sample","transformations","run",".","This","checks","that","unconstrained","values","passed","through","the","event","space","bijectors","into","`","unnormalized_log_prob","`","and","sample","transformations","yield","finite","values",".","This","also","verifies","that","the","transformed","values","have","the","expected","shape",".","Args",":","model",":","The","model","to","validate",".","sample_transformation_shapes",":","Shapes","of","the","transformation","outputs",".","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","ground","truth","standard","error",".","check_ground_truth_standard_deviation",":","Whether","to","check","the","shape","of","the","ground","truth","standard","deviation",".","check_ground_truth_standard_deviation_standard_error",":","Whether","to","check","the","shape","of","the","ground","truth","standard","deviation","standard","error",".","seed",":","Optional","seed","to","use",".","By","default",",","`","test_util.test_seed","(",")","`","is","used","."]
code: def validate_log_prob_and_transforms(
      self,
      model,
      sample_transformation_shapes,
      check_ground_truth_mean=False,
      check_ground_truth_mean_standard_error=False,
      check_ground_truth_standard_deviation=False,
      check_ground_truth_standard_deviation_standard_error=False,
      seed=None,
  ):
      """Validate that the model's log probability and sample transformations run.
  
      This checks that unconstrained values passed through the event space
      bijectors into `unnormalized_log_prob` and sample transformations yield
      finite values. This also verifies that the transformed values have the
      expected shape.
  
      Args:
        model: The model to validate.
        sample_transformation_shapes: Shapes of the transformation outputs.
        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
          ground truth standard error.
        check_ground_truth_standard_deviation: Whether to check the shape of the
          ground truth standard deviation.
        check_ground_truth_standard_deviation_standard_error: Whether to check the
          shape of the ground truth standard deviation standard error.
        seed: Optional seed to use. By default, `test_util.test_seed()` is used.
      """
      batch_size = 16
  
      if seed is None:
          seed = test_util.test_seed(sampler_type="stateless")
  
      def _random_element(
          shape, dtype, default_event_space_bijector, seed
      ):
          unconstrained_shape = (
              default_event_space_bijector.inverse_event_shape(shape)
          )
          unconstrained_shape = tf.TensorShape(
              [batch_size]
          ).concatenate(unconstrained_shape)
          return default_event_space_bijector.forward(
              tf.random.stateless_normal(
                  unconstrained_shape, dtype=dtype, seed=seed
              )
          )
  
      num_seeds = len(tf.nest.flatten(model.dtype))
      flat_seed = tf.unstack(
          tfp.random.split_seed(seed, num_seeds), axis=0
      )
      seed = tf.nest.pack_sequence_as(model.dtype, flat_seed)
  
      test_points = tf.nest.map_structure(
          _random_element,
          model.event_shape,
          model.dtype,
          model.default_event_space_bijector,
          seed,
      )
      log_prob = self.evaluate(model.unnormalized_log_prob(test_points))
  
      self.assertAllFinite(log_prob)
      self.assertEqual((batch_size,), log_prob.shape)
  
      for (
          name,
          sample_transformation,
      ) in model.sample_transformations.items():
          transformed_points = self.evaluate(
              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)
  
          self.assertAllAssertsNested(
              _assertions_part,
              sample_transformation_shapes[name],
              sample_transformation.dtype,
              transformed_points,
              shallow=transformed_points,
              msg="Checking outputs of: {}".format(name),
          )
  
          def _ground_truth_shape_check_part(
              expected_shape, ground_truth
          ):
              self.assertEqual(
                  tuple(expected_shape), tuple(ground_truth.shape)
              )
  
          if check_ground_truth_mean:
              self.assertAllAssertsNested(
                  _ground_truth_shape_check_part,
                  sample_transformation_shapes[name],
                  sample_transformation.ground_truth_mean,
                  shallow=transformed_points,
                  msg="Checking ground truth mean of: {}".format(name),
              )
  
          if check_ground_truth_mean_standard_error:
              self.assertAllAssertsNested(
                  _ground_truth_shape_check_part,
                  sample_transformation_shapes[name],
                  sample_transformation.ground_truth_mean_standard_error,
                  shallow=transformed_points,
                  msg="Checking ground truth mean standard error: {}".format(
                      name
                  ),
              )
  
          if check_ground_truth_standard_deviation:
              self.assertAllAssertsNested(
                  _ground_truth_shape_check_part,
                  sample_transformation_shapes[name],
                  sample_transformation.ground_truth_standard_deviation,
                  shallow=transformed_points,
                  msg="Checking ground truth standard deviation: {}".format(
                      name
                  ),
              )
  
          if check_ground_truth_standard_deviation_standard_error:
              self.assertAllAssertsNested(
                  _ground_truth_shape_check_part,
                  sample_transformation_shapes[name],
                  sample_transformation.ground_truth_standard_deviation_standard_error,
                  shallow=transformed_points,
                  msg="Checking ground truth standard deviation strandard error: {}".format(
                      name
                  ),
              )
  
imports: import collections
  import importlib
  from absl import flags
  from absl import logging
  import jax
  import numpy
  import tensorflow.compat.v2
  import tensorflow_probability
  from tensorflow_probability.python.internal import test_util
label_desc: Use image node_id 1 for calling the InferenceGymTestCase obj's underlying member method code with example usage: obj.validate_log_prob_and_transforms(model, sample_transformation_shapes, check_ground_truth_mean, check_ground_truth_mean_standard_error, check_ground_truth_standard_deviation, check_ground_truth_standard_deviation_standard_error, seed) and returns: default_event_space_bijector