import torch


def create_S_batch(batch_size, device):
    """
    Creates a batch of diag([1,1,0]) matrices
    Input: int
    Output (batch, 3, 3)
    """
    return torch.stack([
        torch.stack([torch.ones(batch_size, device=device),
                     torch.zeros(batch_size, device=device),
                     torch.zeros(batch_size, device=device)],
                    dim=1),
        torch.stack([torch.zeros(batch_size, device=device),
                     torch.ones(batch_size, device=device),
                     torch.zeros(batch_size, device=device)],
                    dim=1),
        torch.stack([torch.zeros(batch_size, device=device),
                     torch.zeros(batch_size, device=device),
                     torch.zeros(batch_size, device=device)],
                    dim=1)
    ], dim=2)