import torch
import math


def sparsify_dense_nre_maps(nre_maps: torch.Tensor, cardinal_omega: float):
    """Sparsify dense NRE maps.
    Args:
        * nre_maps: The [N x H x W] dense NRE maps.
        * cardinal_omega: Card(Omega).
    Returns:
        * sparse_nre_maps: The sparse NRE maps.
        * num_non_zero: The number of non-zero values per map.
    """
    uniform_mask = nre_maps >= math.log(cardinal_omega)
    nre_maps[uniform_mask] = 0.0
    num_non_zero = (~uniform_mask).sum((1, 2))
    return nre_maps.to_sparse(), num_non_zero