Task

class braintools.cogtask.Task(phases=None, input_features=None, output_features=None, trial_init=None, name=None, output_mode='categorical', seed=None, num_classes=None, dt=None, **kwargs)

A cognitive task composed of phases.

The Task class orchestrates phase execution and provides a dataset interface for integration with DataLoaders. Supports both instance-based and class-based definition patterns.

Class-Based Definition

Subclasses can define tasks by overriding class attributes and methods:

• Class attributes: t_fixation, t_sample, t_delay, num_stimuli, etc.

• define_features(): Return (input_features, output_features)

• define_phases(): Return the phase structure

• trial_init(ctx): Initialize trial-level state

Examples

Instance-based (traditional):

>>> import brainunit as u
>>> from braintools.cogtask import Feature, Fixation, Stimulus, Response, circular, label
>>> fix = Feature(1, 'fixation')
>>> stim = Feature(8, 'stimulus')
>>> choice = Feature(2, 'choice')
>>> task = Task(
...     phases=(
...         Fixation(100 * u.ms, inputs={'fixation': 1.0})
...         >> Stimulus(2000 * u.ms, inputs={'stimulus': circular('direction')})
...         >> Response(100 * u.ms, outputs={'label': label('ground_truth')})
...     ),
...     input_features=fix + stim,
...     output_features=fix + choice,
...     trial_init=lambda ctx: ctx.update(
...         ground_truth=ctx.rng.choice(2),
...         direction=ctx.rng.uniform(0, 2 * 3.14159)
...     )
... )

Class-based (new):

>>> class MyTask(Task):
...     t_fixation = 300 * u.ms
...     num_stimuli = 8
...
...     def define_features(self):
...         fix = Feature(1, 'fixation')
...         stim = Feature(self.num_stimuli, 'stimulus')
...         return fix + stim, fix + Feature(2, 'response')
...
...     def define_phases(self):
...         return (Fixation(self.t_fixation, inputs={'fixation': 1.0})
...                 >> Response(100 * u.ms, outputs={'label': label('ground_truth')}))
...
...     def trial_init(self, ctx):
...         ctx['ground_truth'] = ctx.rng.choice(2)
...
>>> task = MyTask(num_stimuli=16, seed=42)

param phases:

The phase structure. If None, uses define_phases() method.

type phases:

Phase

type input_features:

Feature or FeatureSet, optional

param input_features:

Input feature definitions. If None, uses define_features() method.

type input_features:

Feature or FeatureSet, optional

type output_features:

Feature or FeatureSet, optional

param output_features:

Output feature definitions. If None, uses define_features() method.

type output_features:

Feature or FeatureSet, optional

param trial_init:

Function called at the start of each trial. If None and phases is None, uses the trial_init() method.

type trial_init:

Callable[[Context], None] | None

param name:

Task name (defaults to class name).

type name:

str | None

type dt:

float or Quantity, optional

param dt:

Time step used to resolve phase durations into a timestep count. If None (default), the ambient brainstate.environ.get_dt() is used, preserving the previous behaviour. When set, the value is pinned for the whole of trial generation so buffer sizes and the reported dt stay consistent regardless of the ambient environment.

type dt:

float or Quantity, optional

type **kwargs:

param **kwargs:

Override class attributes (e.g., t_fixation=500*u.ms, num_stimuli=16).

batch_sample(size, /, time_first=True, return_meta=False, start_index=0, return_mask=False)

Sample a batch of size trials with indices start_index..start_index+size-1.

When the task was constructed with seed=..., each trial in the batch uses jax.random.fold_in(PRNGKey(seed), start_index + i) so calling batch_sample with the same start_index is reproducible, and successive calls with different start_index produce non-overlapping batches.

Parameters:

return_mask (bool) – If True, also return a (T, B) (or (B, T)) boolean mask of valid timesteps. Required for variable-length tasks if you want to know which trailing positions are padding. The mask is always-True for fixed-length tasks.

define_features()

Define input and output features.

Override in subclass for class-based task definition.

Return type:

Tuple[Any, Any]

Returns:

• input_features (Feature or FeatureSet) – Input feature definitions.

• output_features (Feature or FeatureSet) – Output feature definitions.

define_phases()

Define the phase structure.

Override in subclass for class-based task definition.

Returns:

The task phase structure (single phase or composition).

Return type:

Phase

property dt

Task time step.

Returns the value passed as dt= at construction, or falls back to the ambient brainstate.environ.get_dt() when none was given.

property is_variable_length: bool

True if any phase in the tree has is_variable = True.

Variable-length tasks allocate trial buffers of size max_trial_duration and return a per-timestep mask alongside X/Y from batch_sample() (use return_mask=True).

max_trial_duration(ctx=None)

Static upper bound on the trial’s timestep count.

For fixed tasks this equals the sum of per-phase get_duration outputs. For variable-length tasks each phase contributes its max_steps (e.g. ceil(max_duration / dt) for VariableDuration). The result is a Python int and is safe to use as a static buffer dimension under JIT/vmap.

Return type:

int

property num_classes: int | None

Number of classes for the categorical 1-D label target.

In output_mode='categorical' the target Y is a 1-D integer label array, and the size of a classifier head is the number of distinct label values — which is independent of num_outputs (the summed output-feature dimensionality). num_outputs only coincidentally equals the class count in the default configurations and drifts apart when e.g. cue_dim changes, so prefer num_classes when wiring a categorical model head.

Returns the explicit value passed as num_classes= when set, otherwise falls back to num_outputs (preserving historical behaviour). Returns None in output_mode='vector', where the target is a continuous vector of width num_outputs.

sample_trial(index=0, key=None)

Generate one trial.

Parameters:

• index (int) – Trial index, made available to trial_init as ctx['trial_index']. If the task was constructed with seed=..., the per-trial RNG key is jax.random.fold_in(PRNGKey(seed), index), so reproducibility is keyed on (seed, index).

• key (Array | None) – Explicit PRNG key. Overrides the (seed, index) derivation when given.

trial_init(ctx)

Initialize trial-level state.

Override in subclass to set up trial parameters like ground_truth, stimulus indices, etc.

Parameters:

ctx (Context) – Trial context to populate with state.

Return type:

None