code: def _update_nonzero_block_chol(
      chol, idx, psd_matrix, new_nonzeros, previous_nonzeros
  ):
      """Efficient update to the cholesky factor of the 'slab' (nonzero) submatrix.
  
      This performs an efficient update when `nonzeros` changes by a single entry.
      It is equivalent to
  
      ```
      updated_chol = cholesky(_select_nonzero_block(psd_matrix, new_nonzeros))
      ```
  
      but requires time only quadratic in the dimension `num_features`, where
      a naive recomputation would take cubic time.
  
  
      Args:
        chol: (batch of) float Tensor lower-triangular Cholesky factor(s) of
          `select_nonzero_block(psd_matrix, previous_nonzeros)`, where
          `previous_nonzeros` differs from `new_nonzeros` in the `idx`th entry only.
          This has shape `[num_features, num_features]`.
        idx: scalar int `Tensor` feature index in `[0, num_features)`.
        psd_matrix: (batch of) float Tensor positive semidefinite matrix(s) of shape
          `[num_features, num_features]`.
        new_nonzeros: (batch of) boolean Tensor vectors of shape `[num_features]`.
        previous_nonzeros: (batch of) boolean Tensor vectors of shape
          `[num_features]`.
  
      Returns:
        updated_chol: (batch of) float Tensor lower-triangular Cholesky factor(s) of
          `select_nonzero_block(psd_matrix, new_nonzeros)`.
      """
      psd_row = tf.where(new_nonzeros, psd_matrix[..., idx, :], 0.0)
      eye_row = _set_vector_index(tf.zeros_like(psd_row), idx, 1.0)
      new_row = tf.where(
          new_nonzeros[..., idx, tf.newaxis], psd_row, eye_row
      )
      # NOTE: We could also compute `old_row` from `chol`, but we believe it is
      # more numerically accurate to use `psd_matrix`, as `chol` may have
      # accumulated errors over multiple calls to `_update_nonzero_block_chol`.
      old_row = _select_nonzero_block(psd_matrix, previous_nonzeros)[
          ..., idx, :
      ]
      return _symmetric_increment_chol(
          chol,
          idx=idx,
          # Set the `idx`th row/col to its target value if the `idx`th feature is
          # now nonzero; otherwise set it to the identity.
          increment=new_row - old_row,
      )
  
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 2 for calling a global function with example usage: _update_nonzero_block_chol(chol, idx, psd_matrix, new_nonzeros, previous_nonzeros) and returns: _symmetric_increment_chol