NDTStitchMasker#

class pretrain.models.NDTStitchMasker(mask_ratio, max_block_size=1, block_mask_prob=0.5)[source]#

Bases: torch.nn.modules.module.Module

Temporal masking for spikes (BERT-style).

Randomly masks time bins across all neurons. Each call masks either independent bins or contiguous time spans, drawn per call.

Parameters:

  • mask_ratio – Fraction of time bins to mask.

  • max_block_size – Maximum block size for contiguous masking, 1 for independent per-bin masking.

  • block_mask_prob – Probability of masking a contiguous block rather than random bins, drawn per call and only read when max_block_size is above 1.

generate_block_mask(shape, device)[source]#

Generates a mask where points are grouped into contiguous blocks.

Parameters:

  • shape – Tuple of (Batch size, Time steps)

  • device – Torch device

forward(spikes)[source]#

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.