Linear

class ts1.models.single_session.Linear(bin_size=0.02)

Bases: core.model.BaseModel

Single linear layer mapping binned spike counts to task outputs.

Notation: \(B\) = batch size, \(T_{in}\) = input time bins, \(N\) = units, \(D_{out}\) = task output dim, \(T_{out}\) = output time steps.

configure_readout() must be called before inference; it fixes \(D_{out}\) and the output shape.

1. input_fn(): bin raw spikes into \((T_{in}, N)\) then flatten to \((T_{in} \cdot N,)\).

2. forward(): project \((B, T_{in} \cdot N)\) and reshape to \((B, 1, D_{out})\) or \((B, T_{out}, D_{out})\) depending on the target resolution.

Parameters:

bin_size (float) – Width of each time bin in seconds.

configure_readout(readout_spec)

Fix \(D_{out}\) and build the linear projection head.

The output shape depends on the target resolution:

• Sequence-level (\(T_{out}=1\)): linear \((T_{in} \cdot N) \to D_{out}\), reshaped to \((B, 1, D_{out})\).

• Timestep-level: \(D_{out}\) is expanded by \(T_{out}\), so linear \((T_{in} \cdot N) \to D_{out} \cdot T_{out}\), reshaped to \((B, T_{out}, D_{out})\).

Parameters:

readout_spec (ReadoutSpec) – Task specification carrying \(D_{out}\) and the target resolution.

input_fn(data)

Bin and flatten spikes.

Parameters:

data (Data) – Trial data containing raw spike times and unit metadata.

Return type:

dict[str, Tensor]

Returns:

Dict with model_inputs.spikes of shape \((T_{in} \cdot N,)\).

forward(spikes)

Map flattened binned spikes to task predictions.

Parameters:

spikes (Tensor) – \((B, T_{in} \cdot N)\) flattened binned spike counts.

Return type:

Tensor

Returns:

\((B, 1, D_{out})\) for sequence-level tasks or \((B, T_{out}, D_{out})\) for timestep-level tasks.

classmethod create_search_space(trial, cfg)

Map out the model’s Optuna search space.

Call trial.suggest_*; the names suggested become the keys process_tunable_params() receives.

Parameters:

• trial (Trial) – Optuna trial to register suggestions on.

• cfg (DictConfig) – The run config, for values the space depends on.

classmethod process_tunable_params(tune_params)

Turn suggested hyperparameters into config overrides.

Runs before the config is filled, so this is where a suggestion is mapped onto the config path it sets (batch_size_log2 -> batch_size), a value is derived from another, or a default is supplied for something not being tuned.

Parameters:

tune_params (dict) – The names create_search_space() suggested, with their values.

Return type:

dict

Returns:

The overrides to apply to the config. The default returns them unchanged.