Function test_spike_count_binary_lif_node

W8ASpike/Int2Spike/test.py:12–54

Test function for SpikeCountBinaryLIFNode. Args: x (torch.Tensor, optional): Input tensor with non-negative integer values. If None, a random tensor will be generated. high (int): Upper bound for random generation (exclusive). siz

(
    x: torch.Tensor = None, 
    low: int = 0, 
    high: int = 15, 
    size: tuple = (12, 1024, 2048)
)

)

def test_spike_count_binary_lif_node(
    x: torch.Tensor = None, 
    low: int = 0, 
    high: int = 15, 
    size: tuple = (12, 1024, 2048)
) -> bool:
    """
    Test function for SpikeCountBinaryLIFNode.

    Args:
        x (torch.Tensor, optional): Input tensor with non-negative integer values.
                                    If None, a random tensor will be generated.
        high (int): Upper bound for random generation (exclusive).
        size (tuple): Shape of randomly generated tensor.

    Returns:
        bool: True if the spike sequence sums match the original input; False otherwise.
    """
    if x is None:
        x = torch.randint(low=low, high=high+1, size=size)
    else:
        if not isinstance(x, torch.Tensor):
            raise TypeError("Input x must be a torch.Tensor")

        if not torch.all(x >= 0):
            raise ValueError("Input x must be non-negative")

        if not torch.allclose(x, x.round()):
            raise ValueError("Input x must contain integer values only")
   
    x_zero = torch.zeros_like(x, dtype=torch.float32)
    
    lif = SpikeCountBinaryLIFNode()
    spike_sum = spike_fake_quant(x, lif, x_zero)

    match = torch.allclose((x + x_zero).to(dtype=spike_sum.dtype), spike_sum, rtol=1e-3, atol=1e-3)
    
    if match:
        print(f"\u2714 Numerical match: {lif.__class__.__name__} spike sum matches input spike counts.")
    else:
        print(f"\u2718 Mismatch: {lif.__class__.__name__} spike sum does not match input spike counts.")

    return match

def test_spike_count_ternary_lif_node(
    x: torch.Tensor = None,