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: Precompute quantities needed to sample with given targets.
  
  This method computes a sampler state (including factorized precision
  matrices) from scratch for a given sparsity pattern. This requires
  time proportional to `num_features**3`. If a sampler state is already
  available for an off-by-one sparsity pattern, the `_flip_feature` method
  (which takes time proportional to `num_features**2`) is
  generally more efficient.
  
  Args:
    targets: (batch of) float Tensor regression outputs of shape
      `[num_outputs]`.
    nonzeros: (batch of) boolean Tensor vectors of shape `[num_features]`.
    observation_noise_variance: float Tensor of to scale the posterior
      precision.
  
  Returns:
    sampler_state: instance of `SpikeSlabSamplerState` collecting (potentially
      batched) Tensor quantities relevant to the sampler. See
      `SpikeSlabSamplerState` for details.
code: def _initialize_sampler_state(
      self, targets, nonzeros, observation_noise_variance
  ):
      """Precompute quantities needed to sample with given targets.
  
      This method computes a sampler state (including factorized precision
      matrices) from scratch for a given sparsity pattern. This requires
      time proportional to `num_features**3`. If a sampler state is already
      available for an off-by-one sparsity pattern, the `_flip_feature` method
      (which takes time proportional to `num_features**2`) is
      generally more efficient.
  
      Args:
        targets: (batch of) float Tensor regression outputs of shape
          `[num_outputs]`.
        nonzeros: (batch of) boolean Tensor vectors of shape `[num_features]`.
        observation_noise_variance: float Tensor of to scale the posterior
          precision.
  
      Returns:
        sampler_state: instance of `SpikeSlabSamplerState` collecting (potentially
          batched) Tensor quantities relevant to the sampler. See
          `SpikeSlabSamplerState` for details.
      """
      with tf.name_scope("initialize_sampler_state"):
          targets = tf.convert_to_tensor(targets, dtype=self.dtype)
          nonzeros = tf.convert_to_tensor(nonzeros, dtype=tf.bool)
  
          x_transpose_y = tf.linalg.matvec(
              self.design_matrix, targets, adjoint_a=True
          )
  
          # Ensure that `nonzeros` has full batch shape.
          batch_shape = ps.shape(x_transpose_y)[:-1]
          nonzeros = tf.broadcast_to(
              nonzeros,
              ps.broadcast_shape(
                  ps.shape(nonzeros),
                  ps.concat([batch_shape, [1]], axis=0),
              ),
          )
  
          weights_posterior_precision = (
              self.x_transpose_x
              + self.weights_prior_precision
              * observation_noise_variance
          )
          conditional_prior_precision_chol = tf.linalg.cholesky(
              _select_nonzero_block(
                  self.weights_prior_precision, nonzeros
              )
          )
          conditional_posterior_precision_chol = tf.linalg.cholesky(
              _select_nonzero_block(
                  weights_posterior_precision, nonzeros
              )
          )
          conditional_weights_mean = tf.where(
              nonzeros,
              tf.linalg.cholesky_solve(
                  conditional_posterior_precision_chol,
                  x_transpose_y[..., tf.newaxis],
              )[..., 0],
              0.0,
          )
          return self._compute_log_prob(
              x_transpose_y=x_transpose_y,
              nonzeros=nonzeros,
              conditional_prior_precision_chol=conditional_prior_precision_chol,
              conditional_posterior_precision_chol=conditional_posterior_precision_chol,
              weights_posterior_precision=weights_posterior_precision,
              conditional_weights_mean=conditional_weights_mean,
              observation_noise_variance_posterior_scale=(
                  # SS_gamma / 2 from eqn (7) of [1].
                  self.observation_noise_variance_prior_scale
                  + (  # ss / 2
                      tf.reduce_sum(targets**2, axis=-1)
                      - tf.reduce_sum(  # y'y  # beta_gamma' V_gamma^{-1} beta_gamma
                          conditional_weights_mean * x_transpose_y,
                          axis=-1,
                      )
                  )
                  / 2
              ),
          )
  
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 3 for calling the SpikeSlabSampler obj's underlying member method code with example usage: obj._initialize_sampler_state(targets, nonzeros, observation_noise_variance) and returns: self