code: def test_posterior_on_nonzero_subset_matches_bayesian_regression(
      self, default_pseudo_observations
  ):
      # Generate a synthetic regression task.
      design_matrix, _, targets = self.evaluate(
          self._random_regression_task(
              num_features=5,
              num_outputs=20,
              batch_shape=[2],
              seed=test_util.test_seed(),
          )
      )
  
      # Utilities to extract values for nonzero-weight features.
      nonzeros = np.array([True, False, True, False, True])
      nonzero_subvector = lambda x: x[..., nonzeros]
      nonzero_submatrix = lambda x: self.evaluate(x)[..., nonzeros][
          ..., nonzeros, :
      ]
  
      # Compute the weight posterior mean and precision for these nonzeros.
      sampler = spike_and_slab.SpikeSlabSampler(
          design_matrix,
          default_pseudo_observations=default_pseudo_observations,
      )
      initial_state = sampler._initialize_sampler_state(
          targets=targets,
          nonzeros=nonzeros,
          observation_noise_variance=1.0,
      )
  
      # Compute the analytic posterior for the regression problem restricted to
      # only the selected features. Note that by slicing a submatrix of the
      # prior precision we are implicitly *conditioning* on having observed the
      # other weights to be zero (which is sensible in this case), versus slicing
      # into the covariance which would give the marginal (unconditional) prior
      # on the selected weights.
      (
          restricted_weights_posterior_mean,
          restricted_weights_posterior_prec,
      ) = ncp.mvn_conjugate_linear_update(
          prior_scale=tf.linalg.cholesky(
              tf.linalg.inv(
                  nonzero_submatrix(sampler.weights_prior_precision)
              )
          ),
          linear_transformation=nonzero_subvector(design_matrix),
          likelihood_scale=tf.eye(20),
          observation=targets,
      )
  
      # The sampler's posterior should match the posterior from the restricted
      # problem.
      self.assertAllClose(
          nonzero_subvector(
              self.evaluate(initial_state.conditional_weights_mean)
          ),
          restricted_weights_posterior_mean,
          atol=5e-5,
      )
      self.assertAllClose(
          nonzero_submatrix(
              initial_state.conditional_posterior_precision_chol
          ),
          tf.linalg.cholesky(
              restricted_weights_posterior_prec.to_dense()
          ),
      )
  
imports: from absl.testing import parameterized
  import numpy
  import tensorflow.compat.v2
  from tensorflow_probability.python.distributions import inverse_gamma
  from tensorflow_probability.python.distributions import normal_conjugate_posteriors
  from tensorflow_probability.python.experimental.sts_gibbs import spike_and_slab
  from tensorflow_probability.python.internal import samplers
  from tensorflow_probability.python.internal import test_util
  from tensorflow_probability.python.mcmc.internal import util
label_desc: Use image node_id 5 for calling the SpikeAndSlabTest obj's underlying member method code with example usage: obj.test_posterior_on_nonzero_subset_matches_bayesian_regression(default_pseudo_observations) without return types