from typing import Optional
import torch


def batch_eye(
    eye_dim: int, batch_size: int, device: Optional[torch.device] = None
) -> torch.Tensor:
    """Create a batch of identity matrices.

    Parameters
    ----------
    eye_dim : int
        Desired eye dimension.
    batch_size : int
        Desired batch dimension.
    device : Optional[torch.device], default=None
        The device.

    Returns
    -------
    torch.Tensor, shape=(batch_size, eye_dim, eye_dim)
        Batched eye.
    """
    if device is None:
        eye = torch.eye(eye_dim).unsqueeze(0)  # shape=(1, eye_dim, eye_dim)
    else:
        eye = torch.eye(eye_dim, device=device).unsqueeze(0)
    return eye.expand([batch_size, eye_dim, eye_dim])