class_docstr: Sampler for Bayesian regression with a spike-and-slab prior on weights.
  
  This implementation follows the sampler described in section 3.2
  of Scott and Varian, 2013 [1].
  
  ### Model
  
  This sampler assumes the regression model
  
  ```
  y ~ Normal(loc=matvec(design_matrix,  # `X` in [1].
                        weights),       # `beta` in `[1]`.
             scale=observation_noise_scale)  # `sigma_epsilon` in [1].
  ```
  
  where the design matrix has shape `[num_outputs, num_features]`, with a
  conjugate InverseGamma prior on the noise variance (eqn (6) of [1]):
  
  ```
  observation_noise_scale**2 ~ InverseGamma(
    concentration=observation_noise_variance_prior_concentration,
    scale=observation_noise_variance_prior_scale)
  ```
  
  and a spike-and-slab prior on the weights (eqns (5) and (6) of [1]):
  
  ```
  slab_weights ~ MultivariateNormal(
     loc=0.,  # `b` from [1].
     precision=(weights_prior_precision  # `Omega^{-1}` from [1].
                / observation_noise_scale**2))
  nonzeros ~ Bernoulli(probs=nonzero_prior_prob)  # `gamma` from [1].
  weights = slab_weights * nonzeros
  ```
  
  ### Example
  
  Constructing a sampler instance specifies the model priors:
  
  ```python
  sampler = spike_and_slab.SpikeSlabSampler(
    design_matrix=design_matrix,
    observation_noise_variance_prior_concentration=1.,
    observation_noise_variance_prior_scale=1.
    nonzero_prior_prob=0.1)
  ```
  
  The sampler instance itself is stateless, though some internal methods take
  or accept `SpikeSlabSamplerState` tuples representing posterior quantities
  maintained within a sampling pass. The sampler is
  invoked by passing the regression targets (`y`) and the initial sparsity
  pattern (`nonzeros`):
  
  ```
  (observation_noise_variance,
   weights) = sampler.sample_noise_variance_and_weights(
     targets=y, initial_nonzeros=tf.ones([num_features], dtype=tf.bool))
  ```
  
  This implements the stochastic search variable selection (SSVS) algorithm [2],
  sweeping over the features in random order to resample their sparsity
  indicators one by one. It then returns a sample from the joint posterior
  on the regression weights and the observation noise variance, conditioned
  on the resampled sparsity pattern.
  
  #### References
  
  [1] Steven L. Scott and Hal Varian. Predicting the Present with Bayesian
      Structural Time Series. __International Journal of Mathematical Modelling
      and Numerical Optimisation 5.1-2 (2014): 4-23.__
      https://people.ischool.berkeley.edu/~hal/Papers/2013/pred-present-with-bsts.pdf
  
  [2] George, E. I. and McCulloch, R. E. Approaches for Bayesian variable
      selection. __Statistica Sinica 7, 339–374 (1997)__.
class_docstr_tok: ["Sampler","for","Bayesian","regression","with","a","spike-and-slab","prior","on","weights",".","This","implementation","follows","the","sampler","described","in","section","3.2","of","Scott","and","Varian",",","2013","[","1","]",".","#","#","#","Model","This","sampler","assumes","the","regression","model","``","`","y","~","Normal","(","loc=matvec","(","design_matrix",",","#","`","X","`","in","[","1","]",".","weights",")",",","#","`","beta","`","in","`","[","1","]","`",".","scale=observation_noise_scale",")","#","`","sigma_epsilon","`","in","[","1","]",".","``","`","where","the","design","matrix","has","shape","`","[","num_outputs",",","num_features","]","`",",","with","a","conjugate","InverseGamma","prior","on","the","noise","variance","(","eqn","(","6",")","of","[","1","]",")",":","``","`","observation_noise_scale","*","*","2","~","InverseGamma","(","concentration=observation_noise_variance_prior_concentration",",","scale=observation_noise_variance_prior_scale",")","``","`","and","a","spike-and-slab","prior","on","the","weights","(","eqns","(","5",")","and","(","6",")","of","[","1","]",")",":","``","`","slab_weights","~","MultivariateNormal","(","loc=0.",",","#","`","b","`","from","[","1","]",".","precision=","(","weights_prior_precision","#","`","Omega^","{","-1","}","`","from","[","1","]",".","\/","observation_noise_scale","*","*","2",")",")","nonzeros","~","Bernoulli","(","probs=nonzero_prior_prob",")","#","`","gamma","`","from","[","1","]",".","weights","=","slab_weights","*","nonzeros","``","`","#","#","#","Example","Constructing","a","sampler","instance","specifies","the","model","priors",":","``","`","python","sampler","=","spike_and_slab.SpikeSlabSampler","(","design_matrix=design_matrix",",","observation_noise_variance_prior_concentration=1.",",","observation_noise_variance_prior_scale=1",".","nonzero_prior_prob=0.1",")","``","`","The","sampler","instance","itself","is","stateless",",","though","some","internal","methods","take","or","accept","`","SpikeSlabSamplerState","`","tuples","representing","posterior","quantities","maintained","within","a","sampling","pass",".","The","sampler","is","invoked","by","passing","the","regression","targets","(","`","y","`",")","and","the","initial","sparsity","pattern","(","`","nonzeros","`",")",":","``","`","(","observation_noise_variance",",","weights",")","=","sampler.sample_noise_variance_and_weights","(","targets=y",",","initial_nonzeros=tf.ones","(","[","num_features","]",",","dtype=tf.bool",")",")","``","`","This","implements","the","stochastic","search","variable","selection","(","SSVS",")","algorithm","[","2","]",",","sweeping","over","the","features","in","random","order","to","resample","their","sparsity","indicators","one","by","one",".","It","then","returns","a","sample","from","the","joint","posterior","on","the","regression","weights","and","the","observation","noise","variance",",","conditioned","on","the","resampled","sparsity","pattern",".","#","#","#","#","References","[","1","]","Steven","L.","Scott","and","Hal","Varian",".","Predicting","the","Present","with","Bayesian","Structural","Time","Series",".","__International","Journal","of","Mathematical","Modelling","and","Numerical","Optimisation","5.1-2","(","2014",")",":","4-23.__","https",":","\/\/people.ischool.berkeley.edu\/~hal\/Papers\/2013\/pred-present-with-bsts.pdf","[","2","]","George",",","E.","I.","and","McCulloch",",","R.","E.","Approaches","for","Bayesian","variable","selection",".","__Statistica","Sinica","7",",","339\u2013374","(","1997",")","__","."]
docstr: Computes an unnormalized log prob of a sampler state.
  
  This corresponds to equation (8) in [1]. It scores a sparsity pattern by
  the marginal likelihood of the observed targets (ignoring constant terms
  that do not depend on the sparsity pattern) multiplied by the prior
  probability of the sparsity pattern.
  
  Args: See `SpikeSlabSamplerState`.
  
  Returns:
    sampler_state: a `SpikeSlabSamplerState` instance containing the given
      args and the corresponding unnormalized log prob.
code: def _compute_log_prob(
      self,
      x_transpose_y,
      nonzeros,
      conditional_prior_precision_chol,
      conditional_posterior_precision_chol,
      conditional_weights_mean,
      weights_posterior_precision,
      observation_noise_variance_posterior_scale,
  ):  # pylint: disable=g-doc-args
      """Computes an unnormalized log prob of a sampler state.
  
      This corresponds to equation (8) in [1]. It scores a sparsity pattern by
      the marginal likelihood of the observed targets (ignoring constant terms
      that do not depend on the sparsity pattern) multiplied by the prior
      probability of the sparsity pattern.
  
      Args: See `SpikeSlabSamplerState`.
  
      Returns:
        sampler_state: a `SpikeSlabSamplerState` instance containing the given
          args and the corresponding unnormalized log prob.
      """
      return SpikeSlabSamplerState(
          x_transpose_y=x_transpose_y,
          nonzeros=nonzeros,
          conditional_prior_precision_chol=conditional_prior_precision_chol,
          conditional_posterior_precision_chol=conditional_posterior_precision_chol,
          conditional_weights_mean=conditional_weights_mean,
          weights_posterior_precision=weights_posterior_precision,
          observation_noise_variance_posterior_scale=(
              observation_noise_variance_posterior_scale
          ),
          unnormalized_log_prob=(  # Equation (8) of [1].
              _half_logdet(conditional_prior_precision_chol)
              - _half_logdet(conditional_posterior_precision_chol)
              + self.nonzeros_prior.log_prob(nonzeros)
              - (
                  self.observation_noise_variance_posterior_concentration
                  - 1
              )
              * tf.math.log(
                  2 * observation_noise_variance_posterior_scale
              )
          ),
      )
  
imports: import collections
  import tensorflow.compat.v2
  from tensorflow_probability.python.bijectors import softplus
  from tensorflow_probability.python.distributions import bernoulli
  from tensorflow_probability.python.distributions import inverse_gamma
  from tensorflow_probability.python.distributions import joint_distribution_auto_batched
  from tensorflow_probability.python.distributions import sample
  from tensorflow_probability.python.experimental.distributions.mvn_precision_factor_linop import MultivariateNormalPrecisionFactorLinearOperator
  from tensorflow_probability.python.internal import broadcast_util
  from tensorflow_probability.python.internal import dtype_util
  from tensorflow_probability.python.internal import parameter_properties
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.internal import samplers
  from tensorflow_probability.python.internal import vectorization_util
  from tensorflow_probability.python.math import linalg
  from tensorflow_probability.python.mcmc.internal import util
label_desc: Use image node_id 6 for calling the SpikeSlabSampler obj's underlying member method code with example usage: obj._compute_log_prob(x_transpose_y, nonzeros, conditional_prior_precision_chol, conditional_posterior_precision_chol, conditional_weights_mean, weights_posterior_precision, observation_noise_variance_posterior_scale) and returns: SpikeSlabSamplerState