class_docstr: Posterior predictive in a conjugate Multi-task GP regression model.
code: def _compute_flattened_covariance(self, index_points=None):
      # This is of shape KN x KN, where K is the number of outputs
      # Compute this explicitly via the Schur Complement of the vector kernel.
      # The reason this is written explicitly as opposed to using a GPRM
      # internally for reshaping is there is potential for efficiency gains when
      # `observation_noise_variance = 0.`.
      index_points = self._get_index_points(index_points)
      kxx = self.kernel.matrix_over_all_tasks(
          index_points, index_points
      )
  
      kxz = self.kernel.matrix_over_all_tasks(
          index_points, self.observation_index_points
      ).to_dense()
      if self._observations_is_missing is not None:
          kxz = tf.where(
              _vec(self._observations_is_missing)[..., tf.newaxis, :],
              tf.zeros([], dtype=kxz.dtype),
              kxz,
          )
      if self._observation_scale is not None:
          observation_scale = self._observation_scale
      else:
          observation_scale = _compute_observation_scale(
              self.kernel,
              self.observation_index_points,
              self.cholesky_fn,
              observation_noise_variance=self.observation_noise_variance,
              observations_is_missing=self.observations_is_missing,
          )
  
      cholinv_kzx = observation_scale.solve(kxz, adjoint_arg=True)
      kxz_kzzinv_kzx = tf.linalg.matmul(
          cholinv_kzx, cholinv_kzx, transpose_a=True
      )
  
      flattened_covariance = kxx.to_dense() - kxz_kzzinv_kzx
      if self.predictive_noise_variance is None:
          return flattened_covariance
      broadcast_shape = distribution_util.get_broadcast_shape(
          flattened_covariance,
          self.predictive_noise_variance[..., tf.newaxis, tf.newaxis],
      )
      flattened_covariance = tf.broadcast_to(
          flattened_covariance, broadcast_shape
      )
      return _add_diagonal_shift(
          flattened_covariance, self.predictive_noise_variance
      )
  
imports: import functools
  import tensorflow.compat.v2
  from tensorflow_probability.python.bijectors import softplus
  from tensorflow_probability.python.distributions import cholesky_util
  from tensorflow_probability.python.distributions import distribution
  from tensorflow_probability.python.distributions import mvn_linear_operator
  from tensorflow_probability.python.distributions.internal import stochastic_process_util
  from tensorflow_probability.python.experimental.distributions import multitask_gaussian_process
  from tensorflow_probability.python.experimental.linalg import linear_operator_unitary
  from tensorflow_probability.python.experimental.psd_kernels import multitask_kernel
  from tensorflow_probability.python.internal import batch_shape_lib
  from tensorflow_probability.python.internal import distribution_util
  from tensorflow_probability.python.internal import dtype_util
  from tensorflow_probability.python.internal import nest_util
  from tensorflow_probability.python.internal import parameter_properties
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.internal import reparameterization
  from tensorflow_probability.python.internal import tensor_util
  from tensorflow_probability.python.internal import tensorshape_util
  from tensorflow_probability.python.math.psd_kernels.internal import util
label_desc: Use image node_id 18 for calling the MultiTaskGaussianProcessRegressionModel obj's underlying member method code with example usage: obj._compute_flattened_covariance(index_points) and returns: _add_diagonal_shift, flattened_covariance