keras/src/backend/tensorflow/numpy.py/diag
def diag(x, k=0):
    x = convert_to_tensor(x)
    if len(x.shape) == 1:
        return tf.cond(
            tf.equal(tf.size(x), 0),
            lambda: tf.zeros([builtins.abs(k), builtins.abs(k)], dtype=x.dtype),
            lambda: tf.linalg.diag(x, k=k),
        )
    elif len(x.shape) == 2:
        return diagonal(x, offset=k)
    else:
        raise ValueError(f"`x` must be 1d or 2d. Received: x.shape={x.shape}")