import tensorflow as tf

    """
    Check that K=Identity and K=None give same answer
    """
    chol_from_diag = tf.stack([tf.diag(sqrt_diag[:, i]) for i in range(Datum.N)]) # N x M x M
    s = sqrt_diag if diag else chol_from_diag

