keras/src/backend/torch/numpy.py/identity
def identity(n, dtype=None):
    dtype = to_torch_dtype(dtype or config.floatx())

    # TODO: torch.eye doesn't support bfloat16 with cpu
    if get_device() == "cpu" and dtype == torch.bfloat16:
        return cast(
            torch.eye(n, dtype=to_torch_dtype("float32"), device=get_device()),
            dtype,
        )
    return torch.eye(n, dtype=dtype, device=get_device())