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: Loops over all features to resample their sparsity indicators.
  
  The sampler loops over the features in random order, where each iteration
  updates the `nonzeros` indicator for that particular (single) feature
  weight. This update is a collapsed Gibbs sampling step, i.e., it samples
  from the posterior on the current sparsity indicator given the remaining
  indicators, after marginalizing (collapsing) out the observation noise
  variance and the continuous regression weights under their conjugate priors.
  
  Args:
    initial_sampler_state: instance of `DynamicSpikeSlabSamplerState`
      collecting Tensor quantities relevant to the sampler. See
      `DynamicSpikeSlabSamplerState` for details.
    seed: PRNG seed; see `tfp.random.sanitize_seed` for details.
  
  Returns:
    final sampler_state: instance of `DynamicSpikeSlabSamplerState` in which
      the sparsity indicators for all features have been resampled.
code: def _resample_all_features(self, initial_sampler_state, seed):
      """Loops over all features to resample their sparsity indicators.
  
      The sampler loops over the features in random order, where each iteration
      updates the `nonzeros` indicator for that particular (single) feature
      weight. This update is a collapsed Gibbs sampling step, i.e., it samples
      from the posterior on the current sparsity indicator given the remaining
      indicators, after marginalizing (collapsing) out the observation noise
      variance and the continuous regression weights under their conjugate priors.
  
      Args:
        initial_sampler_state: instance of `DynamicSpikeSlabSamplerState`
          collecting Tensor quantities relevant to the sampler. See
          `DynamicSpikeSlabSamplerState` for details.
        seed: PRNG seed; see `tfp.random.sanitize_seed` for details.
  
      Returns:
        final sampler_state: instance of `DynamicSpikeSlabSamplerState` in which
          the sparsity indicators for all features have been resampled.
      """
      with tf.name_scope("resample_all_features"):
          feature_seed, loop_seed = samplers.split_seed(seed, n=2)
  
          # Visit features in random order.
          feature_permutation = tf.argsort(
              tf.random.stateless_uniform(
                  [self.num_features], seed=feature_seed
              )
          )
  
          @tf.function(autograph=False)
          def resample_one_feature(step, seed, sampler_state):
              seed, next_seed = samplers.split_seed(seed, n=2)
              idx = tf.gather(feature_permutation, step)
  
              # Maybe flip this weight's sparsity indicator.
              proposed_sampler_state = self._flip_feature(
                  sampler_state, idx=idx
              )
              should_flip = bernoulli.Bernoulli(
                  logits=(
                      proposed_sampler_state.unnormalized_log_prob
                      - sampler_state.unnormalized_log_prob
                  ),
                  dtype=tf.bool,
              ).sample(seed=seed)
              return (
                  step + 1,
                  next_seed,
                  mcmc_util.choose(
                      should_flip, proposed_sampler_state, sampler_state
                  ),
              )
  
          _, _, final_sampler_state = tf.while_loop(
              cond=lambda step, *args: step < self.num_features,
              body=resample_one_feature,
              loop_vars=(0, loop_seed, initial_sampler_state),
          )
          return final_sampler_state
  
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 5 for calling the DynamicSpikeSlabSampler obj's underlying member method code with example usage: obj._resample_all_features(initial_sampler_state, seed) and returns: final_sampler_state, unknown,  next_seed,  mcmc_util