keras/src/backend/torch/numpy.py/eye
def eye(N, M=None, k=None, dtype=None):
    dtype = to_torch_dtype(dtype or config.floatx())
    M = N if M is None else M
    k = 0 if k is None else k
    if k == 0:
        # TODO: torch.eye doesn't support bfloat16 with cpu
        if get_device() == "cpu" and dtype == torch.bfloat16:
            return cast(
                torch.eye(
                    N, M, dtype=to_torch_dtype("float32"), device=get_device()
                ),
                dtype,
            )
        return torch.eye(N, M, dtype=dtype, device=get_device())
    diag_length = builtins.max(N, M)
    diag = torch.ones(diag_length, dtype=dtype, device=get_device())
    return torch.diag(diag, diagonal=k)[:N, :M]