Switch

class braintools.cogtask.Switch(selector, cases, default=None, name='Switch')

Multi-way conditional phase selection.

Evaluates the selector function to get a key, then executes the corresponding phase from the cases dictionary.

Examples

>>> # Rule-dependent response
>>> phases = (
...     Stimulus(500 * u.ms)
...     >> Delay(1000 * u.ms)
...     >> Switch(
...         lambda ctx: ctx['rule'],
...         cases={
...             'pro': ProResponse(500 * u.ms),
...             'anti': AntiResponse(500 * u.ms),
...         },
...         default=DefaultResponse(500 * u.ms)
...     )
... )

>>> # Multiple choice selection
>>> phases = Switch(
...     lambda ctx: ctx['choice'],
...     cases={
...         0: Choice0Response(100 * u.ms),
...         1: Choice1Response(100 * u.ms),
...         2: Choice2Response(100 * u.ms),
...     }
... )

Parameters:

• selector (Callable[[Context], Any]) – Function that takes context and returns a key.

• cases (Dict[Any, Phase]) – Mapping from keys to phases.

• default (Phase | None) – Phase to execute if key not found in cases.

• name (str) – Name for this conditional phase.

children()

Return the immediate child phases of a compound phase.

Leaf phases return []. Subclasses like Sequence, Repeat, Parallel, If, Switch, While override this so that the Task can traverse the whole tree to bind features.

encode_inputs(ctx)

Fill ctx.inputs[phase_start:phase_end] with input encoding.

Called once per phase after duration is determined. Must modify ctx.inputs in-place.

Parameters:

ctx (Context) – Context with input buffer and trial state.

Return type:

None

encode_outputs(ctx)

Fill ctx.outputs[phase_start:phase_end] with target encoding.

Called once per phase after duration is determined. Must modify ctx.outputs in-place.

Parameters:

ctx (Context) – Context with output buffer and trial state.

Return type:

None

execute(ctx)

Execute the phase corresponding to the selector’s key.

Return type:

None

execute_packed(ctx)

Packed-mode dispatch.

When the selector returns a concrete Python value (eager sample_trial, or a selector that ignores trial state) we dispatch with an ordinary key in cases lookup, which supports arbitrary hashable keys (e.g. strings).

When the selector returns a tracer (under jit/vmap, i.e. batch_sample) we dispatch with jax.lax.switch() over an ordered list of branches. In this path the case keys must be integers — a traced value cannot be a string. Keys not present in cases (or any value when there is no matching integer key) route to default, or to a no-op when there is no default.

Note

Like If.execute_packed(), the tensor buffers are threaded through the dispatch and are correct, but Python-side ctx state written inside a branch is not reliable under tracing (every branch is traced). Do not rely on branch-written ctx[...] values in get_trial_meta/downstream phases when using batch_sample.

Return type:

None

get_duration(ctx)

Duration depends on which case is selected.

Return type:

int

max_steps(ctx)

Static upper bound on this phase’s length in timesteps.

Must return a Python int with no dependence on traced values. Used by Task in variable-length mode to size shape-stable buffers. The default delegates to get_duration which is correct for fixed-duration phases. Variable-duration phases (e.g. those wrapping TruncExp/UniformDuration) override this to return the truncation upper bound divided by ctx.dt.

Parameters:

ctx (Context) – A stub or trial context providing ctx.dt. The default implementation does not read ctx.rng or trial state.

Returns:

Upper bound on number of timesteps for this phase.

Return type:

int

step_count(ctx)

Traced actual length of this phase in timesteps.

Returns a jax.Array int32 scalar. May depend on ctx[...] values populated by trial_init. Must satisfy 0 <= step_count(ctx) <= max_steps(ctx) for every trial.

The default returns a static value equal to get_duration; that is correct for any phase whose actual length matches its upper bound. Variable-duration phases override this to compute the traced length from ctx state without any int(...) cast.

Return type:

Array