class_docstr: Linear mixing of the latent GPs to form the output.
class_docstr_tok: ["Linear","mixing","of","the","latent","GPs","to","form","the","output","."]
code: def K_diag(
      self, X: TensorType, full_output_cov: bool = True
  ) -> tf.Tensor:
      K = cs(
          tf.stack([k.K_diag(X) for k in self.kernels], axis=-1),
          "[batch..., N, L]",
      )
      rank = tf.rank(X) - 1
      ones = tf.ones((rank,), dtype=tf.int32)
  
      if full_output_cov:
          # Can currently not use einsum due to unknown shape from `tf.stack()`
          # return tf.einsum('nl,lk,lq->nkq', K, self.W, self.W)
          Wt = cs(tf.transpose(self.W), "[L, P]")
          L = tf.shape(Wt)[0]
          P = tf.shape(Wt)[1]
          return cs(
              tf.reduce_sum(
                  cs(K[..., None, None], "[batch..., N, L, 1, 1]")
                  * cs(
                      tf.reshape(Wt, tf.concat([ones, [L, P, 1]], 0)),
                      "[..., L, P, 1]",
                  )
                  * cs(
                      tf.reshape(Wt, tf.concat([ones, [L, 1, P]], 0)),
                      "[..., L, 1, P]",
                  ),
                  axis=-3,
              ),
              "[batch..., N, P, P]",
          )
      else:
          # return tf.einsum('nl,lk,lk->nkq', K, self.W, self.W)
          return cs(
              tf.linalg.matmul(K, self.W**2.0, transpose_b=True),
              "[batch..., N, P]",
          )
  
imports: import abc
  from typing import Optional, Sequence, Tuple
  import tensorflow
  from check_shapes import check_shape
  from check_shapes import check_shapes, inherit_check_shapes
  from ...base import Parameter, TensorType
  from ..base import Combination, Kernel
label_desc: Use image node_id 6 for calling the LinearCoregionalization obj's underlying member method code with example usage: obj.K_diag(X, full_output_cov) and returns: cs, cs