keras/src/ops/numpy.py/identity
def identity(n, dtype=None):
    """Return the identity tensor.

    The identity tensor is a square tensor with ones on the main diagonal and
    zeros elsewhere.

    Args:
        n: Number of rows (and columns) in the `n x n` output tensor.
        dtype: Data type of the output tensor.

    Returns:
        The identity tensor.
    """
    return backend.numpy.identity(n, dtype=dtype)