Replace usages of tfd.MultivariateNormalFullCovariance in TFP.

PiperOrigin-RevId: 464185451

Author: emilyfertig

discussion/neutra/neutra_kernel_test.py

@@ -41,8 +41,8 @@ def testSingleTensor(self, bijector, dtype):
     base_mean = tf.convert_to_tensor(value=[1., 0], dtype=dtype)
     base_cov = tf.convert_to_tensor(value=[[1, 0.5], [0.5, 1]], dtype=dtype)
 
-    base_dist = tfd.MultivariateNormalFullCovariance(
-        loc=base_mean, covariance_matrix=base_cov)
+    base_dist = tfd.MultivariateNormalTriL(
+        loc=base_mean, scale_tril=tf.linalg.cholesky(base_cov))
     target_dist = bijector(base_dist)
 
     def debug_fn(*args):
@@ -90,8 +90,8 @@ def testNested(self, bijector):
     base_mean = tf.constant([1., 0])
     base_cov = tf.constant([[1, 0.5], [0.5, 1]])
 
-    dist_2d = tfd.MultivariateNormalFullCovariance(
-        loc=base_mean, covariance_matrix=base_cov)
+    dist_2d = tfd.MultivariateNormalTriL(
+        loc=base_mean, scale_tril=tf.linalg.cholesky(base_cov))
     dist_4d = tfd.MultivariateNormalDiag(scale_diag=tf.ones(4))
 
     target_dist = tfd.JointDistributionSequential([

tensorflow_probability/examples/jupyter_notebooks/TensorFlow_Distributions_Tutorial.ipynb

@@ -1035,10 +1035,10 @@ def _forward_filter_sequential(self, x, mask=None, final_step_only=False):
                 log_marginal_likelihood=0, timestep=0))
 
       # We could directly construct the batch Distributions
-      # filtered_marginals = tfd.MultivariateNormalFullCovariance(
-      #      filtered_means, filtered_covs)
-      # predicted_marginals = tfd.MultivariateNormalFullCovariance(
-      #      predicted_means, predicted_covs)
+      # filtered_marginals = tfd.MultivariateNormalTriL(
+      #      filtered_means, tf.linalg.cholesky(filtered_covs))
+      # predicted_marginals = tfd.MultivariateNormalTriL(
+      #      predicted_means, tf.linalg.cholesky(predicted_covs))
       # but we choose not to: returning the raw means and covariances
       # saves computation in Eager mode (avoiding an immediate
       # Cholesky factorization that the user may not want) and aids

tensorflow_probability/python/distributions/linear_gaussian_ssm.py

@@ -122,8 +122,9 @@ def test_all_distributions_either_work_or_raise_error(self, dist_name, data):
 
   @test_util.numpy_disable_gradient_test
   def test_multivariate_normal(self):
-    d = tfd.MultivariateNormalFullCovariance(
-        loc=[4., 8.], covariance_matrix=[[11., 0.099], [0.099, 0.1]])
+    d = tfd.MultivariateNormalTriL(
+        loc=[4., 8.],
+        scale_tril=tf.linalg.cholesky([[11., 0.099], [0.099, 0.1]]))
     b = tfp.experimental.bijectors.make_distribution_bijector(d)
     self.assertDistributionIsApproximatelyStandardNormal(tfb.Invert(b)(d))
 

tensorflow_probability/python/experimental/bijectors/distribution_bijectors_test.py

@@ -132,10 +132,8 @@ def target(*args):
 
 
 def assert_mvn_target_conservation(event_size, batch_size, **kwargs):
-  initialization = tfd.MultivariateNormalFullCovariance(
-      loc=tf.zeros(event_size),
-      covariance_matrix=tf.eye(event_size)).sample(
-          batch_size, seed=4)
+  initialization = tfd.MultivariateNormalDiag(
+      loc=tf.zeros(event_size)).sample(batch_size, seed=4)
   samples, leapfrogs = run_nuts_chain(
       event_size, batch_size, num_steps=1,
       initial_state=initialization, **kwargs)

tensorflow_probability/python/experimental/mcmc/nuts_autobatching_test.py

@@ -120,10 +120,8 @@ def run_chain():
 
 def assert_mvn_target_conservation(event_size, batch_size, **kwargs):
   strm = test_util.test_seed_stream()
-  initialization = tfd.MultivariateNormalFullCovariance(
-      loc=tf.zeros(event_size),
-      covariance_matrix=tf.eye(event_size)).sample(
-          batch_size, seed=strm())
+  initialization = tfd.MultivariateNormalDiag(
+      loc=tf.zeros(event_size)).sample(batch_size, seed=strm())
   samples, _ = run_nuts_chain(
       event_size, batch_size, num_steps=1,
       initial_state=initialization, **kwargs)

tensorflow_probability/python/experimental/mcmc/pnuts_test.py

@@ -240,8 +240,8 @@ def test_multivariate_normal_covariance_with_sample_chain(self):
     cov = [[0.36, 0.12, 0.06],
            [0.12, 0.29, -0.13],
            [0.06, -0.13, 0.26]]
-    target = tfp.distributions.MultivariateNormalFullCovariance(
-        loc=mu, covariance_matrix=cov
+    target = tfp.distributions.MultivariateNormalTriL(
+        loc=mu, scale_tril=tf.linalg.cholesky(cov)
     )
     fake_kernel = tfp.mcmc.HamiltonianMonteCarlo(
         target_log_prob_fn=target.log_prob,

tensorflow_probability/python/experimental/mcmc/with_reductions_test.py

@@ -233,21 +233,23 @@ def batch_generator():
            (type(my_filter_results)(1, 2, 3, 2, 3, 2, 3), 2, 1))
 
       # pylint: disable=g-long-lambda,cell-var-from-loop
-      mvn = tfd.MultivariateNormalFullCovariance
       dist = tfd.LinearGaussianStateSpaceModel(
           num_timesteps=nsteps,
           transition_matrix=lambda t: tf.linalg.LinearOperatorFullMatrix(
               tf.gather(transition_matrix, t, axis=0)),
-          transition_noise=lambda t: mvn(
+          transition_noise=lambda t: tfd.MultivariateNormalTriL(
               loc=tf.gather(transition_mean, t, axis=0),
-              covariance_matrix=tf.gather(transition_cov, t, axis=0)),
+              scale_tril=tf.linalg.cholesky(
+                  tf.gather(transition_cov, t, axis=0))),
           observation_matrix=lambda t: tf.linalg.LinearOperatorFullMatrix(
               tf.gather(observation_matrix, t, axis=0)),
-          observation_noise=lambda t: mvn(
+          observation_noise=lambda t: tfd.MultivariateNormalTriL(
               loc=tf.gather(observation_mean, t, axis=0),
-              covariance_matrix=tf.gather(observation_cov, t, axis=0)),
-          initial_state_prior=mvn(loc=initial_mean,
-                                  covariance_matrix=initial_cov),
+              scale_tril=tf.linalg.cholesky(
+                  tf.gather(observation_cov, t, axis=0))),
+          initial_state_prior=tfd.MultivariateNormalTriL(
+              loc=initial_mean,
+              scale_tril=tf.linalg.cholesky(initial_cov)),
           experimental_parallelize=False)  # Compare against sequential filter.
       # pylint: enable=g-long-lambda,cell-var-from-loop
 

tensorflow_probability/python/experimental/parallel_filter/parallel_kalman_filter_test.py

@@ -84,6 +84,8 @@ def _band_part(input, num_lower, num_upper, name=None):  # pylint: disable=redef
 
 def _cholesky_solve(chol, rhs, name=None):  # pylint: disable=unused-argument
   """Scipy cho_solve does not broadcast, so we must do so explicitly."""
+  chol = ops.convert_to_tensor(chol)
+  rhs = ops.convert_to_tensor(rhs)
   if JAX_MODE:  # But JAX uses XLA, which can do a batched solve.
     chol = chol + np.zeros(rhs.shape[:-2] + (1, 1), dtype=chol.dtype)
     rhs = rhs + np.zeros(chol.shape[:-2] + (1, 1), dtype=rhs.dtype)
@@ -368,7 +370,7 @@ def _triangular_solve(matrix, rhs, lower=True, adjoint=False, name=None):  # pyl
 
 cholesky = utils.copy_docstring(
     'tf.linalg.cholesky',
-    lambda input, name=None: np.linalg.cholesky(input))
+    lambda input, name=None: np.linalg.cholesky(ops.convert_to_tensor(input)))
 
 cholesky_solve = utils.copy_docstring(
     'tf.linalg.cholesky_solve',

tensorflow_probability/python/internal/backend/numpy/linalg_impl.py

@@ -116,10 +116,8 @@ def run_chain():
 
 def assert_mvn_target_conservation(event_size, batch_size, **kwargs):
   strm = test_util.test_seed_stream()
-  initialization = tfd.MultivariateNormalFullCovariance(
-      loc=tf.zeros(event_size),
-      covariance_matrix=tf.eye(event_size)).sample(
-          batch_size, seed=strm())
+  initialization = tfd.MultivariateNormalDiag(
+      loc=tf.zeros(event_size)).sample(batch_size, seed=strm())
   samples, _ = run_nuts_chain(
       event_size, batch_size, num_steps=1,
       initial_state=initialization, **kwargs)