Add errors for MTGP/MTGPRM when observations are ill-formed.

PiperOrigin-RevId: 473825902

Author: srvasude

tensorflow_probability/python/experimental/distributions/BUILD

@@ -187,6 +187,8 @@ multi_substrate_py_library(
         "//tensorflow_probability/python/experimental/linalg:linear_operator_unitary",
         "//tensorflow_probability/python/experimental/psd_kernels:multitask_kernel",
         "//tensorflow_probability/python/internal:dtype_util",
+        "//tensorflow_probability/python/internal:tensor_util",
+        "//tensorflow_probability/python/internal:tensorshape_util",
     ],
 )
 
@@ -216,6 +218,8 @@ multi_substrate_py_library(
         # tensorflow dep,
         "//tensorflow_probability/python/distributions:cholesky_util",
         "//tensorflow_probability/python/internal:dtype_util",
+        "//tensorflow_probability/python/internal:tensor_util",
+        "//tensorflow_probability/python/internal:tensorshape_util",
         "//tensorflow_probability/python/math/psd_kernels/internal:util",
     ],
 )

tensorflow_probability/python/experimental/distributions/multitask_gaussian_process.py

@@ -32,6 +32,7 @@
 from tensorflow_probability.python.internal import prefer_static as ps
 from tensorflow_probability.python.internal import reparameterization
 from tensorflow_probability.python.internal import tensor_util
+from tensorflow_probability.python.internal import tensorshape_util
 
 
 def _vec(x):
@@ -447,7 +448,45 @@ def _get_index_points(self, index_points=None):
     return tf.convert_to_tensor(
         index_points if index_points is not None else self._index_points)
 
+  def _check_observations_valid(self, observations, index_points):
+    observation_rank = tensorshape_util.rank(observations.shape)
+
+    if observation_rank is None:
+      return
+
+    if observation_rank >= 1:
+      # Check that the last dimension of observations matches the number of
+      # tasks.
+      num_observations = tf.compat.dimension_value(observations.shape[-1])
+      if (num_observations is not None and
+          num_observations != 1 and
+          num_observations != self.kernel.num_tasks):
+        raise ValueError(
+            f'Expected the number of observations {num_observations} '
+            f'to broadcast / match the number of tasks '
+            f'{self.kernel.num_tasks}')
+
+    if observation_rank >= 2:
+      num_index_points = tf.compat.dimension_value(observations.shape[-2])
+
+      expected_num_index_points = index_points.shape[
+          -(self.kernel.feature_ndims + 1)]
+      if (num_index_points is not None and
+          expected_num_index_points is not None and
+          num_index_points != 1 and
+          num_index_points != expected_num_index_points):
+        raise ValueError(
+            f'Expected number of index points '
+            f'{expected_num_index_points} to broadcast / match the second '
+            f'to last dimension of `observations` {num_index_points}')
+
   def _log_prob(self, value, index_points=None):
+    # Check that observations with at least 2 dimensions have
+    # shape that's broadcastable to `[N, T]`, where `N` is the number
+    # of index points, and T the number of tasks.
+    index_points = self._get_index_points(index_points)
+    self._check_observations_valid(value, index_points)
+
     return self._get_flattened_marginal_distribution(
         index_points=index_points).log_prob(_vec(value))
 

tensorflow_probability/python/experimental/distributions/multitask_gaussian_process_regression_model.py

@@ -286,6 +286,8 @@ def _mean_fn(x):
       self._observations = observations
       self._observations_is_missing = observations_is_missing
 
+      self._check_observations_valid(observations)
+
       if _flattened_conditional_mean_fn is None:
 
         def flattened_conditional_mean_fn(x):
@@ -319,6 +321,38 @@ def flattened_conditional_mean_fn(x):
           parameters=parameters,
           name=name)
 
+  def _check_observations_valid(self, observations):
+    observation_rank = tensorshape_util.rank(observations.shape)
+
+    if observation_rank is None:
+      return
+
+    if observation_rank >= 1:
+      # Check that the last dimension of observations matches the number of
+      # tasks.
+      num_observations = tf.compat.dimension_value(observations.shape[-1])
+      if (num_observations is not None and
+          num_observations != 1 and
+          num_observations != self.kernel.num_tasks):
+        raise ValueError(
+            f'Expected the number of observations {num_observations} '
+            f'to broadcast / match the number of tasks '
+            f'{self.kernel.num_tasks}')
+
+    if observation_rank >= 2:
+      num_index_points = tf.compat.dimension_value(observations.shape[-2])
+
+      expected_num_index_points = self.observation_index_points.shape[
+          -(self.kernel.feature_ndims + 1)]
+      if (num_index_points is not None and
+          expected_num_index_points is not None and
+          num_index_points != 1 and
+          num_index_points != expected_num_index_points):
+        raise ValueError(
+            f'Expected number of observation index points '
+            f'{expected_num_index_points} to broadcast / match the second '
+            f'to last dimension of `observations` {num_index_points}')
+
   @staticmethod
   def precompute_regression_model(
       kernel,

tensorflow_probability/python/experimental/distributions/multitask_gaussian_process_regression_model_test.py

@@ -110,6 +110,47 @@ def testShapes(self, num_tasks):
     self.assertAllEqual(
         self.evaluate(tf.shape(gp.mean())), batch_shape + event_shape)
 
+  def testValidateArgs(self):
+    index_points = np.linspace(-4., 4., 10, dtype=np.float32)
+    index_points = np.reshape(index_points, [5, 2])
+    index_points = np.linspace(-4., 4., 16, dtype=np.float32)
+    observation_index_points = np.reshape(index_points, [8, 2])
+
+    observation_noise_variance = 1e-4
+    kernel = exponentiated_quadratic.ExponentiatedQuadratic()
+    multi_task_kernel = multitask_kernel.Independent(
+        num_tasks=3, base_kernel=kernel)
+    with self.assertRaisesRegexp(ValueError, 'match the number of tasks'):
+      observations = np.linspace(-1., 1., 24).astype(np.float32)
+      mtgprm_lib.MultiTaskGaussianProcessRegressionModel(
+          multi_task_kernel,
+          observation_index_points=observation_index_points,
+          observations=observations,
+          index_points=index_points,
+          observation_noise_variance=observation_noise_variance,
+          validate_args=True)
+
+    with self.assertRaisesRegexp(ValueError, 'match the number of tasks'):
+      observations = np.linspace(-1., 1., 32).reshape(8, 4).astype(np.float32)
+      mtgprm_lib.MultiTaskGaussianProcessRegressionModel(
+          multi_task_kernel,
+          observation_index_points=observation_index_points,
+          observations=observations,
+          index_points=index_points,
+          observation_noise_variance=observation_noise_variance,
+          validate_args=True)
+
+    with self.assertRaisesRegexp(
+        ValueError, 'match the second to last dimension'):
+      observations = np.linspace(-1., 1., 18).reshape(6, 3).astype(np.float32)
+      mtgprm_lib.MultiTaskGaussianProcessRegressionModel(
+          multi_task_kernel,
+          observation_index_points=observation_index_points,
+          observations=observations,
+          index_points=index_points,
+          observation_noise_variance=observation_noise_variance,
+          validate_args=True)
+
   @parameterized.parameters(1, 3, 5)
   def testBindingIndexPoints(self, num_tasks):
     amplitude = np.float64(0.5)

tensorflow_probability/python/experimental/distributions/multitask_gaussian_process_test.py

@@ -382,6 +382,33 @@ def testMultiTaskBlockSeparable(self):
         self.evaluate(actual_multitask_var),
         self.evaluate(multitask_var), rtol=4e-3)
 
+  def testLogProbValidateArgs(self):
+    index_points = np.linspace(-4., 4., 10, dtype=np.float32)
+    index_points = np.reshape(index_points, [-1, 2])
+
+    observation_noise_variance = 1e-4
+    kernel = exponentiated_quadratic.ExponentiatedQuadratic()
+    multi_task_kernel = multitask_kernel.Independent(
+        num_tasks=3, base_kernel=kernel)
+    multitask_gp = multitask_gaussian_process.MultiTaskGaussianProcess(
+        multi_task_kernel,
+        index_points,
+        observation_noise_variance=observation_noise_variance,
+        validate_args=True)
+
+    with self.assertRaisesRegexp(ValueError, 'match the number of tasks'):
+      observations = np.linspace(-1., 1., 15).astype(np.float32)
+      multitask_gp.log_prob(observations)
+
+    with self.assertRaisesRegexp(ValueError, 'match the number of tasks'):
+      observations = np.linspace(-1., 1., 20).reshape(5, 4).astype(np.float32)
+      multitask_gp.log_prob(observations)
+
+    with self.assertRaisesRegexp(
+        ValueError, 'match the second to last dimension'):
+      observations = np.linspace(-1., 1., 18).reshape(6, 3).astype(np.float32)
+      multitask_gp.log_prob(observations)
+
   def testLogProbMatchesGP(self):
     # Check that the independent kernel parameterization matches using a
     # single-task GP.