def diag(x, k=0):
    """Extract a diagonal or construct a diagonal array.

    Args:
        x: Input tensor. If `x` is 2-D, returns the k-th diagonal of `x`.
            If `x` is 1-D, return a 2-D tensor with `x` on the k-th diagonal.
        k: The diagonal to consider. Defaults to `0`. Use `k > 0` for diagonals
            above the main diagonal, and `k < 0` for diagonals below
            the main diagonal.

    Returns:
        The extracted diagonal or constructed diagonal tensor.

    Examples:
    >>> from keras.src import ops
    >>> x = ops.arange(9).reshape((3, 3))
    >>> x
    array([[0, 1, 2],
           [3, 4, 5],
           [6, 7, 8]])

    >>> ops.diag(x)
    array([0, 4, 8])
    >>> ops.diag(x, k=1)
    array([1, 5])
    >>> ops.diag(x, k=-1)
    array([3, 7])

    >>> ops.diag(ops.diag(x)))
    array([[0, 0, 0],
           [0, 4, 0],
           [0, 0, 8]])
    """
    if any_symbolic_tensors((x,)):
        return Diag(k=k).symbolic_call(x)
    return backend.numpy.diag(x, k=k)
Extract a diagonal or construct a diagonal array.

    Args:
        x: Input tensor. If `x` is 2-D, returns the k-th diagonal of `x`.
            If `x` is 1-D, return a 2-D tensor with `x` on the k-th diagonal.
        k: The diagonal to consider. Defaults to `0`. Use `k > 0` for diagonals
            above the main diagonal, and `k < 0` for diagonals below
            the main diagonal.

    Returns:
        The extracted diagonal or constructed diagonal tensor.

    Examples:
    >>> from keras.src import ops
    >>> x = ops.arange(9).reshape((3, 3))
    >>> x
    array([[0, 1, 2],
           [3, 4, 5],
           [6, 7, 8]])

    >>> ops.diag(x)
    array([0, 4, 8])
    >>> ops.diag(x, k=1)
    array([1, 5])
    >>> ops.diag(x, k=-1)
    array([3, 7])

    >>> ops.diag(ops.diag(x)))
    array([[0, 0, 0],
           [0, 4, 0],
           [0, 0, 8]])
    
keras/src/ops/numpy.py