Support MultivariateNormalPrecisionFactorLinearOperator as a weight prior.

PiperOrigin-RevId: 451466503

Author: ColCarroll

tensorflow_probability/python/experimental/sts_gibbs/gibbs_sampler.py

@@ -140,13 +140,16 @@ class SpikeAndSlabSparseLinearRegression(sts_components.LinearRegression):
 
   def __init__(self,
                design_matrix,
-               weights_prior=None,
+               weights_prior,
                sparse_weights_nonzero_prob=0.5,
                name=None):
     # Extract precision matrix from a multivariate normal prior.
     weights_prior_precision = None
     if hasattr(weights_prior, 'precision'):
-      weights_prior_precision = weights_prior.precision()
+      if isinstance(weights_prior.precision, tf.linalg.LinearOperator):
+        weights_prior_precision = weights_prior.precision.to_dense()
+      else:
+        weights_prior_precision = weights_prior.precision()
     elif weights_prior is not None:
       inverse_scale = weights_prior.scale.inverse()
       weights_prior_precision = inverse_scale.matmul(
@@ -840,7 +843,6 @@ def _build_sampler_loop_body(model,
         sampler = dynamic_spike_and_slab.DynamicSpikeSlabSampler
       else:
         sampler = spike_and_slab.SpikeSlabSampler
-
       spike_and_slab_sampler = sampler(
           design_matrix,
           weights_prior_precision=regression_component._weights_prior_precision,  # pylint: disable=protected-access

tensorflow_probability/python/experimental/sts_gibbs/gibbs_sampler_test.py

@@ -29,6 +29,7 @@
 
 
 tfd = tfp.distributions
+tfde = tfp.experimental.distributions
 tfl = tf.linalg
 
 JAX_MODE = False
@@ -398,6 +399,18 @@ def test_invalid_model_spec_raises_error(self):
           level_variance_prior=tfd.InverseGamma(0.01, 0.01),
           observation_noise_variance_prior=tfd.LogNormal(0., 3.))
 
+  def test_model_with_linop_precision_works(self):
+    observed_time_series = tf.ones([2])
+    design_matrix = tf.eye(2)
+    sampler = gibbs_sampler.build_model_for_gibbs_fitting(
+        observed_time_series,
+        design_matrix=design_matrix,
+        weights_prior=tfde.MultivariateNormalPrecisionFactorLinearOperator(
+            precision_factor=tf.linalg.LinearOperatorDiag(tf.ones(2))),
+        level_variance_prior=tfd.InverseGamma(0.01, 0.01),
+        observation_noise_variance_prior=tfd.InverseGamma(0.01, 0.01))
+    self.assertIsNotNone(sampler)
+
   def test_invalid_options_with_none_design_matrix_raises_error(self):
     observed_time_series = tf.ones([2])
     with self.assertRaisesRegex(