code: def test_whitened(diag: bool) -> None:
      """
      Check that K=Identity and K=None give same answer
      """
      chol_from_diag = tf.stack(
          [
              tf.linalg.diag(Datum.sqrt_diag[:, i])
              for i in range(Datum.N)
          ]  # [N, M, M]
      )
      s = Datum.sqrt_diag if diag else chol_from_diag
  
      kl_white = gauss_kl(Datum.mu, s)
      kl_nonwhite = gauss_kl(Datum.mu, s, Datum.I)
  
      np.testing.assert_allclose(kl_white, kl_nonwhite)
  
imports: from unittest.mock import patch
  import numpy
  import pytest
  import tensorflow
  from check_shapes import ShapeChecker, check_shapes
  from numpy.testing import assert_allclose, assert_almost_equal
  import gpflow
  from gpflow import Parameter, default_float, default_jitter
  from gpflow.base import AnyNDArray, TensorType
  from gpflow.inducing_variables import InducingPoints
  from gpflow.kernels import Kernel
  from gpflow.kullback_leiblers import gauss_kl, prior_kl
  from gpflow.utilities.bijectors import triangular
label_desc: Use image node_id 9 for calling a global function with example usage: test_whitened(diag) without return types