MultitaskReadout#

class core.nn.MultitaskReadout(dim, readout_specs, task_index)[source]#

Linear readout with one head per task, routed by integer index.

Parameters:

dim (int) – Dimensionality of the incoming output embeddings.
readout_specs (Mapping[str, ReadoutSpec]) – Task name to ibl_bwb_eval.tasks.ReadoutSpec. Keys name the projections and the entries of the returned dicts.
task_index (Mapping[str, int]) – Task name to the integer written into output_readout_index. Routing uses these, not ReadoutSpec.id, which is a string.

forward(output_embs, output_readout_index, unpack_output=False)[source]#

Project padded output embeddings through their per-task heads.

Parameters:

output_embs (Tensor) – Output embeddings, (batch, n_out, dim).
output_readout_index (Tensor) – Routing index per output token, (batch, n_out).
unpack_output (bool) – Return one dict per batch sample instead of one dict with every sample’s queries concatenated.

Return type:

dict[str, Tensor] | list[dict[str, Tensor]]

Returns:

{task_name: (total_queries, n_channels)}, or a list of such dicts with (n_queries, n_channels) entries when unpack_output is set.

forward_varlen(output_embs, output_readout_index, output_batch_index, unpack_output=False)[source]#

Project chained output embeddings through their per-task heads.

As forward(), but for sequences chained along a single batch dimension rather than padded, which avoids the padding memory.

Parameters:

output_embs (Tensor) – Output embeddings, (total_ntokens, dim).
output_readout_index (Tensor) – Routing index per output token, (total_ntokens,).
output_batch_index (Tensor) – Batch index per output token, (total_ntokens,).
unpack_output (bool) – Return one dict per batch sample. See forward().

Return type:

dict[str, Tensor] | list[dict[str, Tensor]]

Returns:

See forward().