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]. It differs from the SpikeSlabSampler in that
  it selects the submatrix of the design matrix that is active (given by the
  sparsity pattern of the coefficients), before computing a Cholesky
  decomposition. This may provide a speedup by performing a Cholesky on a (much)
  smaller matrix, at the cost of not being able to jit compile the program, or
  to batch (since each batch element may have a different sparsity pattern, so
  the dynamically sized matrices would all be different sizes).
  
  ### 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.DynamicSpikeSlabSampler(
    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 `DynamicSpikeSlabSamplerState` 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)__.
docstr: Initializes priors for the spike and slab sampler.
  
  Args:
    design_matrix: float `Tensor` regression design matrix (`X` in [1]) having
      shape `[num_outputs, num_features]`.
    nonzero_prior_prob: scalar float `Tensor` prior probability of the 'slab',
      i.e., prior probability that any given feature has nonzero weight (`pi`
      in [1]). Default value: `0.5`.
    weights_prior_precision: float `Tensor` complete prior precision matrix
      over the weights, of shape `[num_features, num_features]`. If not
      specified, defaults to the Zellner g-prior specified in `[1]` as
      `Omega^{-1} = kappa * (X'X + diag(X'X)) / (2 * num_outputs)`, in which
      we've plugged in the suggested default of `w = 0.5`. The parameter
      `kappa` is controlled by the `default_pseudo_observations` argument.
      Default value: `None`.
    default_pseudo_observations: scalar float `Tensor` Controls the number of
      pseudo-observations for the prior precision matrix over the weights.
      Corresponds to `kappa` in [1]. See also `weights_prior_precision`.
    observation_noise_variance_prior_concentration: scalar float `Tensor`
      concentration parameter of the inverse gamma prior on the noise
      variance. Corresponds to `nu / 2` in [1]. Default value: 0.005.
    observation_noise_variance_prior_scale: scalar float `Tensor` scale
      parameter of the inverse gamma prior on the noise variance. Corresponds
      to `ss / 2` in [1]. Default value: 0.0025.
    observation_noise_variance_upper_bound: optional scalar float `Tensor`
      maximum value of sampled observation noise variance. Specifying a bound
      can help avoid divergence when the sampler is initialized far from the
      posterior. Default value: `None`.
    num_missing: Optional scalar float `Tensor`. Corrects for how many missing
      values are are coded as zero in the design matrix.
code: def __init__(
      self,
      design_matrix,
      nonzero_prior_prob=0.5,
      weights_prior_precision=None,
      default_pseudo_observations=1.0,
      observation_noise_variance_prior_concentration=0.005,
      observation_noise_variance_prior_scale=0.0025,
      observation_noise_variance_upper_bound=None,
      num_missing=0.0,
  ):
      """Initializes priors for the spike and slab sampler.
  
      Args:
        design_matrix: float `Tensor` regression design matrix (`X` in [1]) having
          shape `[num_outputs, num_features]`.
        nonzero_prior_prob: scalar float `Tensor` prior probability of the 'slab',
          i.e., prior probability that any given feature has nonzero weight (`pi`
          in [1]). Default value: `0.5`.
        weights_prior_precision: float `Tensor` complete prior precision matrix
          over the weights, of shape `[num_features, num_features]`. If not
          specified, defaults to the Zellner g-prior specified in `[1]` as
          `Omega^{-1} = kappa * (X'X + diag(X'X)) / (2 * num_outputs)`, in which
          we've plugged in the suggested default of `w = 0.5`. The parameter
          `kappa` is controlled by the `default_pseudo_observations` argument.
          Default value: `None`.
        default_pseudo_observations: scalar float `Tensor` Controls the number of
          pseudo-observations for the prior precision matrix over the weights.
          Corresponds to `kappa` in [1]. See also `weights_prior_precision`.
        observation_noise_variance_prior_concentration: scalar float `Tensor`
          concentration parameter of the inverse gamma prior on the noise
          variance. Corresponds to `nu / 2` in [1]. Default value: 0.005.
        observation_noise_variance_prior_scale: scalar float `Tensor` scale
          parameter of the inverse gamma prior on the noise variance. Corresponds
          to `ss / 2` in [1]. Default value: 0.0025.
        observation_noise_variance_upper_bound: optional scalar float `Tensor`
          maximum value of sampled observation noise variance. Specifying a bound
          can help avoid divergence when the sampler is initialized far from the
          posterior. Default value: `None`.
        num_missing: Optional scalar float `Tensor`. Corrects for how many missing
          values are are coded as zero in the design matrix.
      """
      with tf.name_scope("spike_slab_sampler"):
          dtype = dtype_util.common_dtype(
              [
                  design_matrix,
                  nonzero_prior_prob,
                  weights_prior_precision,
                  observation_noise_variance_prior_concentration,
                  observation_noise_variance_prior_scale,
                  observation_noise_variance_upper_bound,
                  num_missing,
              ],
              dtype_hint=tf.float32,
          )
          design_matrix = tf.convert_to_tensor(
              design_matrix, dtype=dtype
          )
          nonzero_prior_prob = tf.convert_to_tensor(
              nonzero_prior_prob, dtype=dtype
          )
          observation_noise_variance_prior_concentration = (
              tf.convert_to_tensor(
                  observation_noise_variance_prior_concentration,
                  dtype=dtype,
              )
          )
          observation_noise_variance_prior_scale = tf.convert_to_tensor(
              observation_noise_variance_prior_scale, dtype=dtype
          )
          num_missing = tf.convert_to_tensor(num_missing, dtype=dtype)
          if observation_noise_variance_upper_bound is not None:
              observation_noise_variance_upper_bound = (
                  tf.convert_to_tensor(
                      observation_noise_variance_upper_bound,
                      dtype=dtype,
                  )
              )
  
          design_shape = ps.shape(design_matrix)
          if len(design_shape) > 2:
              raise ValueError(
                  f"DynamicSpikeSlabSampler does not support batched "
                  f"computation, but the design matrix has shape "
                  f"{design_matrix.shape}"
              )
          num_outputs = (
              tf.cast(design_shape[-2], dtype=dtype) - num_missing
          )
          num_features = design_shape[-1]
  
          x_transpose_x = tf.matmul(
              design_matrix, design_matrix, adjoint_a=True
          )
          if weights_prior_precision is None:
              # Default prior: 'Zellner’s g−prior' from section 3.2.1 of [1]:
              #   `omega^{-1} = kappa * (w X'X + (1 − w) diag(X'X))/n`
              # with default `w = 0.5`.
              weights_prior_precision = (
                  default_pseudo_observations
                  * tf.linalg.set_diag(
                      0.5 * x_transpose_x,
                      tf.linalg.diag_part(x_transpose_x),
                  )
                  / num_outputs
              )
  
          observation_noise_variance_posterior_concentration = (
              observation_noise_variance_prior_concentration
              + tf.convert_to_tensor(num_outputs / 2.0, dtype=dtype)
          )
  
          self.num_outputs = num_outputs
          self.num_features = num_features
          self.design_matrix = design_matrix
          self.x_transpose_x = x_transpose_x
          self.dtype = dtype
          self.nonzeros_prior = sample_dist.Sample(
              bernoulli.Bernoulli(probs=nonzero_prior_prob),
              sample_shape=[num_features],
          )
          self.weights_prior_precision = weights_prior_precision
          self.observation_noise_variance_prior_concentration = (
              observation_noise_variance_prior_concentration
          )
          self.observation_noise_variance_prior_scale = (
              observation_noise_variance_prior_scale
          )
          self.observation_noise_variance_upper_bound = (
              observation_noise_variance_upper_bound
          )
          self.observation_noise_variance_posterior_concentration = (
              observation_noise_variance_posterior_concentration
          )
  
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 gamma
  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 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.mcmc.internal import util
label_desc: Use image node_id 1 to create a new DynamicSpikeSlabSampler object  with example: obj = DynamicSpikeSlabSampler(design_matrix, nonzero_prior_prob, weights_prior_precision, default_pseudo_observations, observation_noise_variance_prior_concentration, observation_noise_variance_prior_scale, observation_noise_variance_upper_bound, num_missing)