import torch


def one_hot(type_map: torch.Tensor, n_classes: int) -> torch.Tensor:
    """
    Take in a type map of shape (B, H, W) with class indices as values 
    and reshape it into a tensor of shape (B, C, H, W)

    Args:
    -----------
        type_map (torch.Tensor): 
            type map
        n_classes (int): 
            number of classes in type_map

    Returns:
    -----------
        torch.Tensor: A one hot tensor from the type map of shape 
        (B, C, H, W)
    """
    assert type_map.dtype == torch.int64, (
        f"Wrong type_map dtype: {type_map.dtype}. Shld be torch.int64"
    )

    one_hot = torch.zeros(
        type_map.shape[0], n_classes, *type_map.shape[1:], 
        device=type_map.device, dtype=type_map.dtype
    )

    return one_hot.scatter_(
        dim=1, index=type_map.unsqueeze(1), value=1.0
    ) + 1e-7