model.py

from __future__ import annotations

import torch
from torch import nn
from torch.nn import functional as F


def surrogate_spike(membrane: torch.Tensor, threshold: float = 1.0) -> torch.Tensor:
    hard = (membrane >= threshold).float()
    soft = torch.sigmoid((membrane - threshold) * 10.0)
    return hard + soft - soft.detach()


class LIFSpikingClassifier(nn.Module):
    def __init__(self, hidden_dimensions: int = 64, decay: float = 0.85) -> None:
        super().__init__()
        self.hidden_dimensions = hidden_dimensions
        self.decay = decay
        self.input = nn.Linear(64, hidden_dimensions)
        self.output = nn.Linear(hidden_dimensions, 10)

    def forward(
        self,
        pixels: torch.Tensor,
        *,
        timesteps: int = 24,
        generator: torch.Generator | None = None,
        return_raster: bool = False,
    ) -> tuple[torch.Tensor, torch.Tensor] | tuple[
        torch.Tensor, torch.Tensor, torch.Tensor
    ]:
        membrane = torch.zeros(
            len(pixels),
            self.hidden_dimensions,
            device=pixels.device,
        )
        logits = torch.zeros(len(pixels), 10, device=pixels.device)
        spike_total = torch.zeros((), device=pixels.device)
        input_raster = []
        for _ in range(timesteps):
            random_values = torch.rand(
                pixels.shape,
                generator=generator,
                device=pixels.device,
            )
            input_spikes = (random_values < pixels).float()
            membrane = self.decay * membrane + self.input(input_spikes)
            hidden_spikes = surrogate_spike(membrane)
            membrane = membrane - hidden_spikes.detach()
            logits = logits + self.output(hidden_spikes)
            spike_total = spike_total + hidden_spikes.sum()
            if return_raster:
                input_raster.append(input_spikes)
        spike_rate = spike_total / (len(pixels) * self.hidden_dimensions * timesteps)
        if return_raster:
            return logits / timesteps, spike_rate, torch.stack(input_raster, dim=1)
        return logits / timesteps, spike_rate


class MatchedDenseClassifier(nn.Module):
    def __init__(self, hidden_dimensions: int = 64) -> None:
        super().__init__()
        self.first = nn.Linear(64, hidden_dimensions)
        self.output = nn.Linear(hidden_dimensions, 10)

    def forward(self, pixels: torch.Tensor) -> torch.Tensor:
        return self.output(F.gelu(self.first(pixels)))


def parameter_count(model: nn.Module) -> int:
    return sum(parameter.numel() for parameter in model.parameters())