code: def _select_nonzero_block(matrix, nonzeros):
      """Replaces the `i`th row & col with the identity if not `nonzeros[i]`.
  
      This function effectively selects the 'slab' rows (corresponding to
      features with nonzero weight) of a prior or posterior precision matrix. To
      guarantee static shapes (and to support batching over different sparsity
      patterns), it returns a matrix of the same shape as the input, in which the
      rows not selected are replaced with the corresponding row of an identity
      matrix. The result can thus be viewed as a permutation of a block-diagonal
      matrix:
  
      ```
      | matrix[nonzeros, nonzeros]    0 |
      | 0                             I |
      ```
  
      and this view justifies the correctness of matrix operations such as
      Cholesky factorization on the selected submatrix, although in actuality
      the features are left at their original indices (not permuted).
  
      Args:
        matrix: (batch of) float Tensor matrix(s) of shape `[num_features,
          num_features]`.
        nonzeros: (batch of) boolean Tensor vectors of shape `[num_features]`.
  
      Returns:
        block_matrix: (batch of) float Tensor matrix(s) of the same shape as
          `matrix`, in which `block_matrix[i, j] = matrix[i, j] if
          `(nonzeros[i] and nonzeros[j]) else eye(num_features)[i, j]`. This is
          positive semidefinite if `matrix` is, since it is a permutation
          of a matrix with PSD blocks.
      """
      # Zero out all entries in the not-selected rows.
      masked = tf.where(
          nonzeros[..., tf.newaxis],
          tf.where(nonzeros[..., tf.newaxis, :], matrix, 0.0),
          0.0,
      )
      # Restore a value of 1 on the diagonal of the not-selected rows. This avoids
      # numerical issues by ensuring that the matrix still has full rank.
      return tf.linalg.set_diag(
          masked, tf.where(nonzeros, tf.linalg.diag_part(masked), 1.0)
      )
  
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 1 for calling a global function with example usage: _select_nonzero_block(matrix, nonzeros) and returns: tf