thrindex.snn

​

Import

import thrindex.snn as snn

​

snn.Sequential

model = snn.Sequential(*layers)

• Sequential is compatible with torch.nn.Sequential for all standard PyTorch operations: .parameters(), .state_dict(), .to(device).
• The first layer in a Sequential determines the model’s input dimension, which is stored in the .thx artifact.

​

snn.Dense

layer = snn.Dense(
    in_features: int,
    out_features: int,
    bias: bool = True,
    delays: bool = False,
    max_delay: int = 15,
)

​

snn.LIF

layer = snn.LIF(
    tau_mem: float,
    threshold: float = 1.0,
    tau_syn: float | None = None,
    reset: str = "subtract",
    dt: float = 1.0,
)

• Input: [batch, T, in_features]
• Output: [batch, T, n_neurons] — binary spike raster (values are exactly 0.0 or 1.0 in the forward pass; gradients are surrogate-smoothed in the backward pass)

• n_neurons is determined by the out_features of the preceding Dense layer. There is no n_neurons argument — the LIF layer infers its size from the input shape on the first forward call.
• The membrane state is initialized to zero at the start of each sample. It is not persistent across samples in a batch.

​

snn.Conv2d

layer = snn.Conv2d(
    in_channels: int,
    out_channels: int,
    kernel_size: int | tuple[int, int],
    stride: int = 1,
    padding: int = 0,
    bias: bool = True,
)

• Input: [batch, T, C_in, H, W]
• Output: [batch, T, C_out, H_out, W_out]

​

snn.Flatten

layer = snn.Flatten()

• Input: [batch, T, C, H, W]
• Output: [batch, T, C × H × W]

​

Forward pass

x = input_spikes  # [batch, T, in_features]
for layer in self.layers:
    x = layer(x)
# x is now [batch, T, out_features]