NDT2

class pretrain.models.NDT2(is_ssl, hidden_dim, units_per_patch, max_spikes, max_num_units, bin_size, tokenize_session, tokenize_subject, enc_depth, enc_heads, enc_ffn_mult, dec_depth, dec_heads, dec_ffn_mult, dropout, activation='gelu', pre_norm=True, is_causal=False, finetune_enable=False)

Bases: core.model.BaseModel

Multi-context masked autoencoder over patched spike tokens [Ye et al., 2023].

Reference implementation: context_general_bci.

Not tested at any context_length other than the default (1.0); see Variable context.

link_datasets(train_dataset, val_dataset, test_dataset=None)

Build whatever the model sizes from the datasets.

Called before configure_readout(), so what is read here sizes the readout.

Parameters:

• train_dataset (IBLBrainWideBench2026) – Training split, the one to size from.

• val_dataset (IBLBrainWideBench2026) – Validation split.

• test_dataset (Optional[IBLBrainWideBench2026]) – Test split, or None when the run scores none.

configure_readout(readout_spec)

Size the readout head for the task being evaluated.

Called after link_datasets(). The spec carries dim and num_timesteps, so a head is sized without branching on it.

Parameters:

readout_spec (ReadoutSpec) – See ReadoutSpec.

load_ckpt(ckpt)

Copy pretrained weights out of a checkpoint into this model.

Read only the weights: the trainer restores optimizer and epoch state itself.

Parameters:

ckpt (dict) – The loaded checkpoint, as written by the pretraining run.

input_fn(data)

Convert one trial’s data into the model’s inputs.

Runs per item on the dataloader workers, as a dataset transform, so it returns tensors rather than batches.

Parameters:

data (Data) – One trial, exposing the interval key configure_readout() set.

Return type:

dict[str, Any]

Returns:

Dict whose model_inputs entry is splatted into forward(). Extra top-level keys are the model’s own to read in its trainer.

forward(in_patches, in_time_idx, in_space_idx, in_not_pad, ssl_mask=None, query_idx=None, query_time_idx=None, query_space_idx=None, query_mask=None, session_idx=None, subject_idx=None, task_idx=None)

Run the model on one collated batch.

Takes whatever input_fn() put under model_inputs, splatted in as keyword arguments, and returns what the trainer’s loss reads.

Return type:

dict[str, Tensor]