import torch


def batch_colout(X: torch.Tensor, p_row: float, p_col: float) -> torch.Tensor:
    """Applies ColOut augmentation to a batch of images, dropping the same
    random rows and columns from all images in a batch.

    Args:
        X: Batch of images of shape (N, C, H, W).
        p_row: Fraction of rows to drop (drop along H).
        p_col: Fraction of columns to drop (drop along W).

    Returns:
        torch.Tensor: Input batch tensor with randomly dropped columns and rows.
    """

    # Get the dimensions of the image
    row_size = X.shape[2]
    col_size = X.shape[3]

    # Determine how many rows and columns to keep
    kept_row_size = int((1 - p_row) * row_size)
    kept_col_size = int((1 - p_col) * col_size)

    # Randomly choose indices to keep. Must be sorted for slicing
    kept_row_idx = sorted(torch.randperm(row_size)[:kept_row_size].numpy())
    kept_col_idx = sorted(torch.randperm(col_size)[:kept_col_size].numpy())

    # Keep only the selected row and columns
    X_colout = X[:, :, kept_row_idx, :]
    X_colout = X_colout[:, :, :, kept_col_idx]
    return X_colout