NDTStitch

class pretrain.models.NDTStitch(hidden_dim, bin_size, encoder_num_layers, encoder_dim_feedforward, encoder_num_heads, encoder_dropout, encoder_activation, pre_encoder_dropout, post_encoder_dropout, use_custom_init=False, tfixup_scale_base=None, tfixup_v_scale_factor=None, finetune_enable=False)

Bases: core.model.BaseModel

Multi-session NDT with per-session stitchers [Ye and Pandarinath, 2021].

Reference implementation: neural-data-transformers.

Note: NDTStitch intentionally has a lot of overlap with the single session NDT implementation (i.e., the DRY code principle is not respected). This design choice is motivated by the desire to make an individual model fully understandable for a user without having to navigate between models.

We document/highlight the key differences between single sesion and it’s stitch version.

Removed: * The spike embedding strategy. * attn_mask usage.

Added: * The NDT in/out_stitcher: For each session of the dataset, an NDT stitcher maps dim_1 to dim_2.

• in_stitcher: num_units -> hidden_dim

• out_stitcher: hidden_dim -> num_units

Each stitcher is a single linear layer. A two-layer MLP was also available, and is not scalable: its hidden width grows with num_units, so the parameter count is quadratic in unit count and paid once per session.

• Session embeddings to account for across-session variability.

Updated: * The encoder/transformer is initialized in the __init__() method and no more in link_datasets() as hidden_dim is a parameter.

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.

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 – One trial, exposing the interval key configure_readout() set.

Returns:

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

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.

forward(spikes, positions, session_tokens)

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.

custom_init()

T-Fixup init for the encoder; see core.nn.tfixup_init_().