Module neural_decoders 

Core compute kernels for POYO+, POSSM, NDT3, and CEBRA decoders.

These are the hot-path operations that benefit from Rust acceleration: spike tokenisation, sinusoidal position encoding, scaled dot-product attention, diagonal SSM step, and InfoNCE contrastive loss.

Functions

gaussian_attention

Gaussian attention. Li et al. (2025), scKGBERT.

infonce_loss

InfoNCE contrastive loss. van den Oord et al. (2018); CEBRA.

scaled_dot_product_attention

Scaled dot-product attention.

sinusoidal_position_encode

Sinusoidal position encoding. Vaswani et al. (2017).

ssm_step_diagonal

Diagonal SSM step. Gu et al. (2022), S4D; Ryoo et al. (2025), POSSM.

tokenise_spikes

Convert binary spike trains to sorted (unit_id, timestamp) tokens.

Type Aliases

SpikeToken

Token: (unit_id, timestamp_ms).