From 2ed8a123a81f32e00170e32480d2060da2a1351a Mon Sep 17 00:00:00 2001
From: Kartik Pradeepan <kartikspradeepan@gmail.com>
Date: Wed, 4 Nov 2026 10:53:57 -0500
Subject: [PATCH 1/6] feat(tasks): show several parts of one kind in a trial
 and record each image

TaskContext.present accepts named parts ('image:word_1') so one trial
can show a sequence before a single answer. Prompts carry part names,
and a multi-image trial's recording keeps one step per image, labelled
by part.
---
 brainscore_core/_model/behavioral.py | 38 ++++++++++++++++++----------
 brainscore_core/contract.py          | 10 +++++---
 brainscore_core/responses.py         | 28 ++++++++++++++------
 tests/test_task_responses.py         | 14 ++++++++++
 tests/test_visual_forms.py           | 35 +++++++++++++++++++++++++
 5 files changed, 101 insertions(+), 24 deletions(-)

diff --git a/brainscore_core/_model/behavioral.py b/brainscore_core/_model/behavioral.py
index ad5dde8e..6486c631 100644
--- a/brainscore_core/_model/behavioral.py
+++ b/brainscore_core/_model/behavioral.py
@@ -387,6 +387,14 @@ def call(row, prompt):
                     extra = self._probabilities_after(stimuli, prompts, labels)
                     outputs['behavior:probabilities'] = extra
         outputs['behavior'] = outputs[response.channels[0]]
+        if prompts:
+            # A multi-image trial records one step per image: name each by its part.
+            image_names = [name for (kind, _), name in zip(prompts[0].parts, prompts[0].names)
+                           if kind == 'image']
+            neural = {region: (assembly.assign_coords(part_name=('part', image_names))
+                               if 'part' in assembly.dims and assembly.sizes['part'] == len(image_names)
+                               else assembly)
+                      for region, assembly in neural.items()}
         stamp = self._stamp(path, 'prompt' if path == 'generation' else 'readout extraction')
         return ({c: outputs[c] for c in channels},
                 {r: a.assign_attrs(recording=stamp) for r, a in neural.items()})
@@ -702,7 +710,12 @@ def _trials_to_assembly(trials, stimuli, layer: str, region: str):
     from brainscore_core.supported_data_standards.brainio.assemblies import NeuroidAssembly
 
     rows = []
-    timed = 'time_bin' in trials[0].dims
+    # A clip keeps its time steps and a multi-image trial its parts: one extra dimension.
+    extra = [dim for dim in trials[0].dims if dim not in ('presentation', 'neuroid')]
+    if len(extra) > 1:
+        raise ValueError(f'A recording during a task may have one extra dimension; got {extra}.')
+    timed = bool(extra)
+    extra = extra[0] if extra else None
     for index, trial in enumerate(trials):
         selected = trial.isel(neuroid=trial['layer'].values == layer)
         if selected.sizes['neuroid'] == 0:
@@ -710,11 +723,10 @@ def _trials_to_assembly(trials, stimuli, layer: str, region: str):
                 f'No activity reached layer {layer!r} while answering stimulus '
                 f'{stimuli["stimulus_id"].iloc[index]!r}; that input may not use it.')
         if timed:
-            # A clip keeps its time steps: (time_bin, neuroid) per presentation.
-            if selected.sizes['time_bin'] != trials[0].sizes['time_bin']:
-                raise ValueError('Clips of different lengths gave different numbers of time steps; '
-                                 'present clips of one length to record them together.')
-            rows.append(np.asarray(selected.transpose('presentation', 'time_bin', 'neuroid').values)[0])
+            if selected.sizes[extra] != trials[0].sizes[extra]:
+                raise ValueError(f'Trials gave different numbers of {extra} steps; present trials of one '
+                                 'length (clips of one duration, sequences of one length) together.')
+            rows.append(np.asarray(selected.transpose('presentation', extra, 'neuroid').values)[0])
         else:
             rows.append(np.asarray(selected.values).reshape(-1))
     first = trials[0].isel(neuroid=trials[0]['layer'].values == layer)
@@ -728,14 +740,14 @@ def _trials_to_assembly(trials, stimuli, layer: str, region: str):
         coords[column] = ('presentation', list(stimuli[column].values))
     if not timed:
         return NeuroidAssembly(np.stack(rows), coords=coords, dims=['presentation', 'neuroid'])
-    time_index = first.indexes.get('time_bin')
-    time_names = list(time_index.names) if isinstance(time_index, pd.MultiIndex) else []
-    time_names += [name for name, coord in first.coords.items()
-                   if coord.dims == ('time_bin',) and name != 'time_bin' and name not in time_names]
-    for name in time_names:
-        coords[name] = ('time_bin', np.asarray(first[name].values))
+    extra_index = first.indexes.get(extra)
+    extra_names = list(extra_index.names) if isinstance(extra_index, pd.MultiIndex) else []
+    extra_names += [name for name, coord in first.coords.items()
+                    if coord.dims == (extra,) and name != extra and name not in extra_names]
+    for name in extra_names:
+        coords[name] = (extra, np.asarray(first[name].values))
     return NeuroidAssembly(np.stack(rows), coords=coords,
-                           dims=['presentation', 'time_bin', 'neuroid'])
+                           dims=['presentation', extra, 'neuroid'])
 
 
 # Keep existing pickle references and public type names stable.
diff --git a/brainscore_core/contract.py b/brainscore_core/contract.py
index f4ae3c9d..ae632297 100644
--- a/brainscore_core/contract.py
+++ b/brainscore_core/contract.py
@@ -46,7 +46,9 @@ class TaskContext:
     ``'probabilities'`` keep their meaning on the plain ``behavior`` channel.
 
     ``present`` gives the order of a trial's parts when a stimulus has several
-    (e.g. ``('text', 'video')``: a sentence, then a clip). A replying model is
+    (e.g. ``('text', 'video')``: a sentence, then a clip). Name columns to show
+    several parts of one kind in sequence: ``('image:instruction', 'image:word_1',
+    'image:word_2')``. A replying model is
     shown the instruction, those parts in order, then the response's directive.
 
     ``readout_fit`` says which network a readout is fitted on while a change
@@ -81,9 +83,11 @@ def __post_init__(self):
         elif self.task_type is None:
             raise ValueError("Give a response (Choice, Scale, Text) or task_type='passive'.")
         if self.present is not None:
-            unknown = set(self.present) - {'text', 'image', 'video', 'audio'}
+            kinds = {'text', 'image', 'video', 'audio'}
+            unknown = [entry for entry in self.present if entry.split(':', 1)[0] not in kinds]
             if unknown or isinstance(self.present, str):
-                raise ValueError(f"present lists parts from 'text', 'image', 'video', 'audio'; got {self.present!r}")
+                raise ValueError("present lists parts as a kind ('text', 'image', 'video', 'audio') "
+                                 f"or 'kind:column' (e.g. 'image:word_1'); got {self.present!r}")
             self.present = tuple(self.present)
 
 
diff --git a/brainscore_core/responses.py b/brainscore_core/responses.py
index 427948a1..294aaa27 100644
--- a/brainscore_core/responses.py
+++ b/brainscore_core/responses.py
@@ -126,6 +126,7 @@ class Prompt:
     """
     parts: Tuple[Tuple[str, Any], ...]
     stimulus_id: Optional[str] = None
+    names: Tuple[str, ...] = ()   # one per part: 'instruction', a column name, or 'directive'
 
     def text(self) -> str:
         """The text parts joined, for a model that reads text only."""
@@ -145,17 +146,28 @@ def trial_prompt(task, stimulus_row) -> Prompt:
     if task.present is None and len(present) > 1:
         raise ValueError(f'Each stimulus has several parts {sorted(present)}; say the order people '
                          'saw them in with TaskContext(present=(...)), e.g. present=("text", "video").')
-    missing = [kind for kind in order if kind not in present]
-    if missing:
-        raise ValueError(f'TaskContext.present names {missing}, but the stimulus has only {sorted(present)}.')
-    parts = []
+    shown = []          # (kind, value, name)
+    for entry in order:
+        if ':' in entry:
+            # A named part: 'image:word_1' shows the image in column word_1.
+            kind, column = entry.split(':', 1)
+            value = _cell(stimulus_row, column)
+            if value is None:
+                raise ValueError(f'TaskContext.present names column {column!r}, which this stimulus lacks.')
+            shown.append((kind, value, column))
+        elif entry in present:
+            shown.append((entry, present[entry], entry))
+        else:
+            raise ValueError(f'TaskContext.present names {entry!r}, but the stimulus has only {sorted(present)}.')
+    parts, names = [], []
     if (task.instruction or '').strip():
-        parts.append(('text', task.instruction.strip()))
-    parts += [(kind, present[kind]) for kind in order]
+        parts.append(('text', task.instruction.strip())); names.append('instruction')
+    for kind, value, name in shown:
+        parts.append((kind, value)); names.append(name)
     directive = task.response.directive() if task.response is not None else ''
     if directive:
-        parts.append(('text', directive))
-    return Prompt(tuple(parts), stimulus_id=_cell(stimulus_row, 'stimulus_id'))
+        parts.append(('text', directive)); names.append('directive')
+    return Prompt(tuple(parts), stimulus_id=_cell(stimulus_row, 'stimulus_id'), names=tuple(names))
 
 
 def _cell(row, column):
diff --git a/tests/test_task_responses.py b/tests/test_task_responses.py
index aef0639d..663af060 100644
--- a/tests/test_task_responses.py
+++ b/tests/test_task_responses.py
@@ -167,3 +167,17 @@ def test_an_unordered_multi_part_stimulus_fails_before_any_reply():
     with pytest.raises(ValueError, match='present'):
         model.participate(session)
     assert calls == []
+
+
+def test_named_parts_show_a_sequence_in_order():
+    from brainscore_core.responses import trial_prompt
+    row = {'stimulus_id': 't1', 'instruction': 'rules.png', 'word_1': 'w1.png', 'word_2': 'w2.png'}
+    task = TaskContext(response=Choice(['yes', 'no']),
+                       present=('image:instruction', 'image:word_1', 'image:word_2'))
+    prompt = trial_prompt(task, row)
+    assert prompt.parts[:3] == (('image', 'rules.png'), ('image', 'w1.png'), ('image', 'w2.png'))
+    assert prompt.names == ('instruction', 'word_1', 'word_2', 'directive')
+    with pytest.raises(ValueError, match='word_3'):
+        trial_prompt(TaskContext(response=Choice(['yes', 'no']), present=('image:word_3',)), row)
+    with pytest.raises(ValueError, match='kind'):
+        TaskContext(response=Choice(['yes', 'no']), present=('picture:word_1',))
diff --git a/tests/test_visual_forms.py b/tests/test_visual_forms.py
index 5d15b6b5..48332a94 100644
--- a/tests/test_visual_forms.py
+++ b/tests/test_visual_forms.py
@@ -98,3 +98,38 @@ def test_a_clip_recorded_during_a_task_uses_the_clip_extractor_and_keeps_time():
     assert recorded.dims == ('presentation', 'time_bin', 'neuroid')
     np.testing.assert_array_equal(recorded.values[0, :, 0], [5.0, 6.0])
     assert list(recorded['time_bin_start_ms'].values) == [0.0, 1000.0]
+
+
+def test_a_multi_image_trial_is_recorded_per_image():
+    from contextlib import contextmanager
+    from brainscore_core.model_interface import TaskContext
+    from brainscore_core.responses import Choice
+    from brainscore_core.streaming import BufferedSession
+
+    class PerImage:
+        def __call__(self, stimuli, layers, **kwargs):
+            raise AssertionError('recorded during the reply')
+
+        @contextmanager
+        def record(self, layers):
+            class Take:
+                @staticmethod
+                def take():
+                    values = np.arange(6, dtype=float).reshape(1, 3, 2)   # 3 images, 2 units
+                    return NeuroidAssembly(values, dims=['presentation', 'part', 'neuroid'], coords={
+                        'stimulus_id': ('presentation', ['?']), 'part_index': ('part', [0, 1, 2]),
+                        'neuroid_id': ('neuroid', ['n0', 'n1']), 'layer': ('neuroid', [layers[0]] * 2)})
+            yield Take()
+
+    model = BrainScoreModel('m', extractor={'vision': PerImage()},
+                            region_layer_map={'IT': LayerSelector('blocks.1', modality='vision')},
+                            reply=lambda prompt: 'yes')
+    trials = StimulusSet({'stimulus_id': ['t1'], 'rules': ['r.png'], 'word_1': ['a.png'], 'word_2': ['b.png']})
+    trials.identifier = 'sequence'
+    task = TaskContext(response=Choice(['yes', 'no']), present=('image:rules', 'image:word_1', 'image:word_2'))
+    session = BufferedSession(stimulus_set=trials, task_context=task,
+                              requested_output_channels=['behavior:choice', 'neural:IT'])
+    model.participate(session)
+    recorded = session.collect('neural:IT')
+    assert recorded.dims == ('presentation', 'part', 'neuroid')
+    assert list(recorded['part_name'].values) == ['rules', 'word_1', 'word_2']

From 03270c70010ef872d23d7b0ef761f84abdc00883 Mon Sep 17 00:00:00 2001
From: Kartik Pradeepan <kartikspradeepan@gmail.com>
Date: Wed, 4 Nov 2026 11:09:08 -0500
Subject: [PATCH 2/6] feat(tasks): record regions on every task path

Recording during a task now works during likelihood scoring, next to a
reply with implied probabilities (scored first from the same prompt),
for composite regions, in streaming sessions (outputs sent per trial),
and for extractors without record() through forward hooks on the model
passed as model=.
---
 brainscore_core/_model/behavioral.py | 220 +++++++++++++++++----------
 brainscore_core/_sessions/drivers.py |  39 ++---
 brainscore_core/brainscore_model.py  |   3 +-
 tests/test_recording_everywhere.py   | 147 ++++++++++++++++++
 4 files changed, 306 insertions(+), 103 deletions(-)
 create mode 100644 tests/test_recording_everywhere.py

diff --git a/brainscore_core/_model/behavioral.py b/brainscore_core/_model/behavioral.py
index 6486c631..1ae95d47 100644
--- a/brainscore_core/_model/behavioral.py
+++ b/brainscore_core/_model/behavioral.py
@@ -1,7 +1,7 @@
 """Behavioral readout and generation support for BrainScoreModel."""
 
 from contextlib import ExitStack, contextmanager, nullcontext
-from typing import Any, Callable, Dict, Optional, Sequence, Tuple
+from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple
 
 from ..contract import TaskContext
 
@@ -316,21 +316,16 @@ def _start_response_task(self, task_context) -> None:
     def respond(self, stimuli, channels, regions=()):
         """Answer a response task; return ``({channel: assembly}, {region: NeuroidAssembly})``.
 
-        Each stimulus is presented once on the answer's path. Regions are
-        recorded during that presentation: the prompt pass of a reply, or the
-        readout's own extraction.
+        Each stimulus is presented once on the answer's path, and regions are
+        recorded during that presentation: the prompt pass of a reply or of
+        likelihood scoring, or the readout's own extraction. Implied
+        probabilities next to a reply are scored first, from the same prompt,
+        so the recording is still the prompt pass.
         """
         from ..responses import Choice, Scale, Text, labels_of, trial_prompt
         owner, task = self.owner, self.task_context
         response, path = task.response, self.response_path
-        layers = {region: self._region_layer(region) for region in regions}
-        if layers and path == 'likelihood':
-            raise NotImplementedError('Recording during likelihood scoring is not supported; '
-                                      "use prefer_path='generation' or 'readout'.")
-        if layers and path == 'generation' and self.probability_source is not None:
-            raise NotImplementedError('Implied probabilities need extra forward passes, so they '
-                                      'cannot be recorded in one presentation with a generated '
-                                      "answer; use prefer_path='readout' or drop probabilities.")
+        layers = {region: self._region_layers(region) for region in regions}
         labels = labels_of(response)
         # Every trial's prompt is built first, so a bad layout fails before any compute.
         prompts = [trial_prompt(task, row) for _, row in stimuli.iterrows()]
@@ -347,17 +342,24 @@ def respond(self, stimuli, channels, regions=()):
                 if self.probability_source == 'readout':
                     outputs['behavior:probabilities'] = _implied(values, 'readout')
         else:
+            score_first = path == 'generation' and self.probability_source == 'likelihood'
             if path == 'likelihood':
                 def call(row, prompt):
                     return owner._likelihood_fn(prompt, list(labels))
             elif owner._reply_fn is not None:
                 def call(row, prompt):
-                    return owner._reply_fn(prompt)
+                    # Scoring first keeps the recorded first pass the prompt alone.
+                    scores = owner._likelihood_fn(prompt, list(labels)) if score_first else None
+                    return owner._reply_fn(prompt), scores
             else:   # a registration's answer: it phrases and parses a Choice itself
                 def call(row, prompt):
                     return owner._generation_fn(stimulus_row=row, instruction=task.instruction,
                                                 label_set=list(labels))
             replies, neural = self._per_stimulus(stimuli, layers, call, prompts)
+            scores = None
+            if path == 'generation' and owner._reply_fn is not None:
+                scores = [s for _, s in replies]
+                replies = [r for r, _ in replies]
             if path == 'likelihood':
                 probabilities = _softmax_rows(replies, labels)
                 chosen = [labels[i] for i in probabilities.argmax(1)]
@@ -383,9 +385,12 @@ def call(row, prompt):
                     outputs['behavior:confidence'] = _numbers(
                         [response.parse_confidence(reply) for reply in replies], stimuli,
                         'reply', raw=replies)
-                if self.probability_source is not None:
-                    extra = self._probabilities_after(stimuli, prompts, labels)
-                    outputs['behavior:probabilities'] = extra
+                if score_first:
+                    outputs['behavior:probabilities'] = _implied_array(
+                        _softmax_rows(scores, labels), stimuli, labels, 'likelihood')
+                elif self.probability_source == 'readout':
+                    outputs['behavior:probabilities'] = _implied(self.readout_classifier.predict_proba(
+                        owner._extract_behavioral_features(stimuli)), 'readout')
         outputs['behavior'] = outputs[response.channels[0]]
         if prompts:
             # A multi-image trial records one step per image: name each by its part.
@@ -395,50 +400,74 @@ def call(row, prompt):
                                if 'part' in assembly.dims and assembly.sizes['part'] == len(image_names)
                                else assembly)
                       for region, assembly in neural.items()}
-        stamp = self._stamp(path, 'prompt' if path == 'generation' else 'readout extraction')
+        stamp = self._stamp(path, 'readout extraction' if path == 'readout' else 'prompt')
         return ({c: outputs[c] for c in channels},
                 {r: a.assign_attrs(recording=stamp) for r, a in neural.items()})
 
-    def _probabilities_after(self, stimuli, prompts, labels):
-        """Implied probabilities alongside a generated answer (extra passes, not recorded)."""
-        owner = self.owner
-        if self.probability_source == 'likelihood':
-            scores = [owner._likelihood_fn(prompt, list(labels)) for prompt in prompts]
-            return _implied_array(_softmax_rows(scores, labels), stimuli, labels, 'likelihood')
-        return _implied(self.readout_classifier.predict_proba(
-            self.owner._extract_behavioral_features(stimuli)), 'readout')
+    def _open_recordings(self, stimuli, layers):
+        """One live recording per input type, before any compute.
 
-    def _per_stimulus(self, stimuli, layers, call, prompts=None):
-        """Call ``call(row, prompt)`` once per stimulus, recording ``layers`` on each first pass."""
-        from ..io_catalog import canonical_modality
+        An extractor's own ``record(layers)`` when it has one; otherwise forward
+        hooks on the model passed as ``BrainScoreModel(model=...)``, with layer
+        paths read from that model.
+        """
         owner = self.owner
+        from .dispatch import extractor_for
         by_modality: Dict[str, list] = {}
-        for region, layer in layers.items():
-            by_modality.setdefault(self._region_modality(region), []).append(layer)
+        for region, region_layers in layers.items():
+            by_modality.setdefault(self._region_modality(region), []).extend(region_layers)
         recordings = {}
         for modality, modality_layers in by_modality.items():
-            from .dispatch import extractor_for
+            modality_layers = list(dict.fromkeys(modality_layers))
             extractor = extractor_for(owner, modality, stimuli)
             if extractor is None:
                 raise NotImplementedError(f'No {modality!r} extractor to record from.')
             try:
                 record = extractor.record
             except AttributeError:
+                record = None
+            if record is not None:
+                recordings[modality] = record(modality_layers)
+            elif owner._model is not None:
+                recordings[modality] = _hook_recording(owner._model, modality_layers, owner.identifier)
+            else:
                 raise NotImplementedError(
-                    f'The {modality!r} extractor cannot record while the model answers: '
-                    'it needs a record(layers) method. Use a readout task, or a '
-                    'neural-only session.') from None
-            recordings[modality] = record(list(dict.fromkeys(modality_layers)))
+                    f'The {modality!r} extractor cannot record while the model answers: give it '
+                    'a record(layers) method, or pass the network as BrainScoreModel(model=...).')
+        return recordings
+
+    def _per_stimulus(self, stimuli, layers, call, prompts=None):
+        """Call ``call(row, prompt)`` once per stimulus, recording ``layers`` on each first pass."""
+        recordings = self._open_recordings(stimuli, layers)
         replies, trials = [], []
         with ExitStack() as stack:
             live = {m: stack.enter_context(r) for m, r in recordings.items()}
             for index, (_, row) in enumerate(stimuli.iterrows()):
                 replies.append(call(row, None if prompts is None else prompts[index]))
                 trials.append({m: recording.take() for m, recording in live.items()})
-        neural = {region: _trials_to_assembly([t[self._region_modality(region)] for t in trials],
-                                              stimuli, layer, region)
-                  for region, layer in layers.items()}
-        return replies, neural
+        return replies, self._trial_regions(trials, stimuli, layers)
+
+    def _trial_regions(self, trials, stimuli, layers):
+        """Per-trial recordings -> one NeuroidAssembly per region (composite regions gathered)."""
+        import xarray as xr
+        neural = {}
+        for region, region_layers in layers.items():
+            modality = self._region_modality(region)
+            per_layer = [_trials_to_assembly([trial[modality] for trial in trials], stimuli, layer, region)
+                         for layer in region_layers]
+            neural[region] = self._gather(region, per_layer)
+        return neural
+
+    def _gather(self, region, per_layer):
+        """A region's units: one layer as is; a composite region's units gathered from its layers."""
+        import xarray as xr
+        from ..selection import CompositeSelector
+        selector = self.owner._region_layer_selectors.get(region)
+        if not isinstance(selector, CompositeSelector):
+            return per_layer[0]
+        recorder = self.owner._recorder
+        combined = xr.concat([recorder.ensure_unit_indices(a) for a in per_layer], dim='neuroid')
+        return recorder.gather_composite_units(combined, selector)
 
     @contextmanager
     def suspended(self):
@@ -459,27 +488,22 @@ def respond_and_record(self, stimuli, regions: Sequence[str]):
 
         Returns ``(BehavioralAssembly, {region: NeuroidAssembly})``. With
         ``answer``, the recording is the first forward pass of each answer (the
-        prompt, before any generated token), taken through each extractor's
-        ``record(layers)``. With a readout, the region layers are extracted in
-        the same call as the readout layer. Unsupported requests fail before
-        any compute.
+        prompt, before any generated token). With a readout, the region layers
+        are extracted in the same call as the readout layer.
         """
-        layers = {region: self._region_layer(region) for region in regions}
+        layers = {region: self._region_layers(region) for region in regions}
         if self.use_generation_for_task:
             return self._generate_and_record(stimuli, layers)
         return self._read_out_and_record(stimuli, layers)
 
-    def _region_layer(self, region: str) -> str:
+    def _region_layers(self, region: str) -> List[str]:
         from ..selection import CompositeSelector
         selector = self.owner._region_layer_selectors.get(region)
         if selector is None:
             raise ValueError(f'Region {region!r} is not in region_layer_map.')
         if isinstance(selector, CompositeSelector):
-            # ponytail: composite regions need per-unit gathering; add when a study asks.
-            raise NotImplementedError(
-                f'Region {region!r} combines several layers; record it with a '
-                'neural-only session.')
-        return selector.layer_path
+            return list(selector.layer_paths)
+        return [selector.layer_path]
 
     def _region_modality(self, region: str) -> str:
         from ..io_catalog import canonical_modality
@@ -493,39 +517,17 @@ def _region_modality(self, region: str) -> str:
                 "LayerSelector(name=..., modality=...) to record it during a task.")
         return next(iter(modalities))
 
-    def _generate_and_record(self, stimuli, layers: Dict[str, str]):
-        from ..io_catalog import canonical_modality
-        owner = self.owner
-        by_modality: Dict[str, list] = {}
-        for region, layer in layers.items():
-            by_modality.setdefault(self._region_modality(region), []).append(layer)
-        recordings = {}
-        for modality, modality_layers in by_modality.items():
-            from .dispatch import extractor_for
-            extractor = extractor_for(owner, modality, stimuli)
-            if extractor is None:
-                raise NotImplementedError(f'No {modality!r} extractor to record from.')
-            try:
-                record = extractor.record
-            except AttributeError:
-                raise NotImplementedError(
-                    f'The {modality!r} extractor cannot record while the model answers: '
-                    'it needs a record(layers) method. Use a readout task, or a '
-                    'neural-only session.') from None
-            recordings[modality] = record(list(dict.fromkeys(modality_layers)))
+    def _generate_and_record(self, stimuli, layers):
+        recordings = self._open_recordings(stimuli, layers)
         trials = []
         with ExitStack() as stack:
             live = {m: stack.enter_context(r) for m, r in recordings.items()}
             behavior = self.generate_predictions(
                 stimuli, after_each=lambda: trials.append(
                     {m: recording.take() for m, recording in live.items()}))
-        neural = {}
-        for region, layer in layers.items():
-            modality = self._region_modality(region)
-            neural[region] = _trials_to_assembly(
-                [trial[modality] for trial in trials], stimuli, layer, region
-            ).assign_attrs(recording=self._stamp('answer', 'prompt'))
-        return behavior, neural
+        stamp = self._stamp('answer', 'prompt')
+        return behavior, {region: assembly.assign_attrs(recording=stamp)
+                          for region, assembly in self._trial_regions(trials, stimuli, layers).items()}
 
     def _stamp(self, path, forward_pass):
         """How a recording was made, so task and passive recordings stay distinct."""
@@ -534,7 +536,7 @@ def _stamp(self, path, forward_pass):
                 if self.task_context.response is not None else None,
                 'forward_pass': forward_pass, 'readout_fit': self.task_context.readout_fit}
 
-    def _read_out_and_record(self, stimuli, layers: Dict[str, str], raw: bool = False):
+    def _read_out_and_record(self, stimuli, layers, raw: bool = False):
         from ..responses import Scale
         owner = self.owner
         readout_layer = owner._behavioral_readout_layer
@@ -545,7 +547,7 @@ def _read_out_and_record(self, stimuli, layers: Dict[str, str], raw: bool = Fals
                     f'Region {region!r} reads {self._region_modality(region)!r} input, but '
                     f'the readout presents {modality!r} input; one presentation cannot '
                     'record it.')
-        all_layers = list(dict.fromkeys([readout_layer, *layers.values()]))
+        all_layers = list(dict.fromkeys([readout_layer, *(l for ls in layers.values() for l in ls)]))
         extracted = owner._recorder.ensure_layer_provenance(
             owner._extract_for_modality(stimuli, modality, all_layers), all_layers,
             description='Readout extraction')
@@ -558,11 +560,11 @@ def _read_out_and_record(self, stimuli, layers: Dict[str, str], raw: bool = Fals
         else:
             behavior = self.readout_classifier.predict_proba(features)
         neural = {}
-        for region, layer in layers.items():
-            selected = extracted.isel(neuroid=layer_of == layer)
+        for region, region_layers in layers.items():
+            per_layer = [extracted.isel(neuroid=layer_of == layer) for layer in region_layers]
+            selected = self._gather(region, per_layer)
             neural[region] = selected.assign_coords(
-                region=('neuroid', [region] * selected.sizes['neuroid'])
-            )
+                region=('neuroid', [region] * selected.sizes['neuroid']))
             if not raw:
                 neural[region] = neural[region].assign_attrs(
                     recording=self._stamp('readout', 'readout extraction'))
@@ -638,6 +640,58 @@ def generate_predictions(self, stimuli, after_each: Optional[Callable[[], None]]
         )
 
 
+@contextmanager
+def _hook_recording(model, layers, identifier):
+    """Forward hooks on ``model`` for extractors without ``record()``.
+
+    ``take()`` returns the first pass since the previous take, each layer's
+    output flattened per presentation, as a still-image extractor packages it.
+    """
+    from brainscore_core.supported_data_standards.brainio.assemblies import NeuroidAssembly
+
+    def module_at(path):
+        module = model
+        for part in path.split('.'):
+            module = module[int(part)] if part.isdigit() else getattr(module, part, None)
+            if module is None:
+                raise NotImplementedError(f'Layer {path!r} is not a module of the model passed as model=; '
+                                          'give the extractor a record(layers) method.')
+        return module
+
+    modules = {layer: module_at(layer) for layer in layers}   # unknown paths fail before compute
+    captured = {}
+
+    def keep(layer):
+        def hook(_module, _inputs, output):
+            if layer not in captured:
+                output = output[0] if isinstance(output, (tuple, list)) else output
+                captured[layer] = np.asarray(output.detach().float().cpu().numpy())
+        return hook
+
+    handles = [module.register_forward_hook(keep(layer)) for layer, module in modules.items()]
+
+    class Recording:
+        @staticmethod
+        def take():
+            missing = [layer for layer in layers if layer not in captured]
+            if missing:
+                raise ValueError(f'No forward pass reached {missing} since the last presentation.')
+            rows = {layer: captured[layer].reshape(1, -1) for layer in layers}
+            captured.clear()
+            values = np.concatenate([rows[layer] for layer in layers], axis=1)
+            neuroid_layers = [layer for layer in layers for _ in range(rows[layer].shape[1])]
+            numbers = [i for layer in layers for i in range(rows[layer].shape[1])]
+            return NeuroidAssembly(values, dims=['presentation', 'neuroid'], coords={
+                'stimulus_id': ('presentation', ['?']),
+                'neuroid_id': ('neuroid', [f'{identifier}.{l}.{i}' for l, i in zip(neuroid_layers, numbers)]),
+                'neuroid_num': ('neuroid', numbers), 'layer': ('neuroid', neuroid_layers)})
+    try:
+        yield Recording()
+    finally:
+        for handle in handles:
+            handle.remove()
+
+
 def _presentation(stimuli):
     return {column: ('presentation', list(stimuli[column].values)) for column in stimuli.columns}
 
diff --git a/brainscore_core/_sessions/drivers.py b/brainscore_core/_sessions/drivers.py
index c2b1c141..78f2f574 100644
--- a/brainscore_core/_sessions/drivers.py
+++ b/brainscore_core/_sessions/drivers.py
@@ -394,10 +394,6 @@ def drive_behavior_and_neural_session(model, session, driver: str = "interact")
     task_context = getattr(session, "task_context", None)
     if task_context is None:
         raise ValueError("A behavior session needs session.task_context.")
-    if getattr(session, "streaming", False):
-        raise NotImplementedError(
-            "Recording during a task presents the stimulus set at once; "
-            "use streaming=False.")
     time_bins = getattr(session, "time_bins", None)
     if time_bins is not None and len(time_bins) > 1:
         raise NotImplementedError(
@@ -410,21 +406,26 @@ def drive_behavior_and_neural_session(model, session, driver: str = "interact")
     stimuli = _behavior_stimuli_from_session(session, task_context)
     if stimuli is None:
         raise ValueError("Recording during a task needs stimuli to present.")
-    if task_context.response is not None:
-        behavior, neural = model._behavioral.respond(stimuli, answers, regions)
-    else:
-        if set(answers) - {"behavior"}:
-            raise ValueError(f"{answers} need TaskContext(response=...); task_type "
-                             f"{task_context.task_type!r} answers on 'behavior'.")
-        plain, neural = model._behavioral.respond_and_record(stimuli, regions)
-        behavior = {"behavior": plain} if answers else {}
-    for channel in answers:
-        session.emit(StreamEvent(channel=channel, payload=behavior[channel], t_ms=0.0,
-                                 meta={"driver": driver, "task_type": task_context.task_type}))
-    for region in regions:
-        session.emit(StreamEvent(channel=f"neural:{region}", payload=neural[region], t_ms=0.0,
-                                 meta={"driver": driver, "record": region,
-                                       "during": "behavior"}))
+    if task_context.response is None and set(answers) - {"behavior"}:
+        raise ValueError(f"{answers} need TaskContext(response=...); task_type "
+                         f"{task_context.task_type!r} answers on 'behavior'.")
+    # A streaming session gets each trial's outputs as soon as that trial ends;
+    # otherwise the whole set is answered and sent at once.
+    streaming = getattr(session, "streaming", False)
+    batches = ([(index, _preserve_stimulus_set_type(stimuli, stimuli.iloc[[index]]))
+                for index in range(len(stimuli))] if streaming else [(0, stimuli)])
+    for index, batch in batches:
+        if task_context.response is not None:
+            behavior, neural = model._behavioral.respond(batch, answers, regions)
+        else:
+            plain, neural = model._behavioral.respond_and_record(batch, regions)
+            behavior = {"behavior": plain} if answers else {}
+        for channel in answers:
+            session.emit(StreamEvent(channel=channel, payload=behavior[channel], t_ms=float(index),
+                                     meta={"driver": driver, "task_type": task_context.task_type}))
+        for region in regions:
+            session.emit(StreamEvent(channel=f"neural:{region}", payload=neural[region], t_ms=float(index),
+                                     meta={"driver": driver, "record": region, "during": "behavior"}))
 
 
 _drive_neural_session_via_process = drive_neural_session
diff --git a/brainscore_core/brainscore_model.py b/brainscore_core/brainscore_model.py
index a41646c2..7b43fb6b 100644
--- a/brainscore_core/brainscore_model.py
+++ b/brainscore_core/brainscore_model.py
@@ -139,7 +139,8 @@ def __init__(
             and written answers.
         :param likelihood: ``(prompt, candidates) -> {candidate: log
             probability}``; scores candidate replies without generating. Gives
-            choices and their implied probabilities.
+            choices and their implied probabilities. Run the prompt on its own
+            before any candidate, so a recording during scoring is the prompt pass.
         :param policy: ``EnvironmentStep -> EnvironmentResponse``; chooses an
             action for closed-loop / embodied evaluation.
         :param intervention: ``StateChange -> (PerturbationApplied, cleanup)``;
diff --git a/tests/test_recording_everywhere.py b/tests/test_recording_everywhere.py
new file mode 100644
index 00000000..4dd3df27
--- /dev/null
+++ b/tests/test_recording_everywhere.py
@@ -0,0 +1,147 @@
+"""Regions record on every task path: likelihood scoring, a reply with implied
+probabilities, composite regions, streaming sessions, and extractors without
+record() (through forward hooks on the model)."""
+from contextlib import contextmanager
+
+import numpy as np
+import pytest
+
+from brainscore_core.brainscore_model import BrainScoreModel
+from brainscore_core.model_interface import TaskContext
+from brainscore_core.responses import Choice
+from brainscore_core.selection import CompositeSelector, LayerSelector
+from brainscore_core.streaming import BufferedSession
+from brainscore_core.supported_data_standards.brainio.assemblies import NeuroidAssembly
+from brainscore_core.supported_data_standards.brainio.stimuli import StimulusSet
+
+pytestmark = pytest.mark.unit
+
+STIMULI = StimulusSet({'stimulus_id': ['a', 'b', 'c'], 'sentence': ['x', 'yy', 'zzz']})
+STIMULI.identifier = 'everywhere'
+
+
+class Net:
+    """Each pass sees a text; 'early' has 3 units, 'late' 2. Passes are logged."""
+
+    def __init__(self):
+        self.passes, self.hooks = [], []
+
+    def forward(self, text):
+        self.passes.append(text)
+        out = {'early': np.array([len(text), 2.0 * len(text), 3.0 * len(text)]),
+               'late': np.array([-len(text), 10.0])}
+        for hook in self.hooks:
+            hook(out)
+        return out
+
+
+class Recordable:
+    identifier = backbone_id = 'toy'
+
+    def __init__(self, net):
+        self.net = net
+
+    def __call__(self, stimuli, layers, **kwargs):
+        raise AssertionError('recorded during the task')
+
+    @contextmanager
+    def record(self, layers):
+        seen = []
+
+        class Take:
+            @staticmethod
+            def take():
+                first = seen[0]
+                seen.clear()
+                values = np.concatenate([first[layer] for layer in layers])[None]
+                units = [(layer, i) for layer in layers for i in range(len(first[layer]))]
+                return NeuroidAssembly(values, dims=['presentation', 'neuroid'], coords={
+                    'stimulus_id': ('presentation', ['?']),
+                    'neuroid_id': ('neuroid', [f'{l}.{i}' for l, i in units]),
+                    'neuroid_num': ('neuroid', [i for _, i in units]),
+                    'layer': ('neuroid', [l for l, _ in units])})
+        self.net.hooks.append(seen.append)
+        try:
+            yield Take()
+        finally:
+            self.net.hooks.remove(seen.append)
+
+
+def _model(net, region_map=None):
+    def reply(prompt):
+        net.forward(prompt.text())
+        return 'cat'
+
+    def likelihood(prompt, candidates):
+        net.forward(prompt.text())                         # the prompt alone first
+        return {c: -float(len(net.forward(prompt.text() + c)['late'])) - i
+                for i, c in enumerate(candidates)}
+    return BrainScoreModel('toy', extractor={'text': Recordable(net)},
+                           region_layer_map=region_map or {'A': LayerSelector('early', modality='text')},
+                           reply=reply, likelihood=likelihood)
+
+
+def _run(model, task, channels, streaming=False):
+    session = BufferedSession(stimulus_set=STIMULI, task_context=task,
+                              requested_output_channels=channels, streaming=streaming)
+    model.participate(session)
+    return session
+
+
+def _prompt_text(stimulus_text):
+    return f'Which?\n\n{stimulus_text}\n\nAnswer with exactly one of: cat, dog.'
+
+
+def test_recording_during_likelihood_scoring_is_the_prompt_pass():
+    net = Net()
+    task = TaskContext(instruction='Which?', response=Choice(['cat', 'dog'], probabilities=True),
+                       prefer_path='likelihood')
+    recorded = _run(_model(net), task, ['behavior:probabilities', 'neural:A']).collect('neural:A')
+    expected = [Net().forward(_prompt_text(s))['early'] for s in ('x', 'yy', 'zzz')]
+    np.testing.assert_array_equal(recorded.values, expected)
+
+
+def test_a_reply_with_implied_probabilities_records_the_prompt_pass():
+    net = Net()
+    task = TaskContext(instruction='Which?', response=Choice(['cat', 'dog'], probabilities=True),
+                       prefer_path='generation')
+    session = _run(_model(net), task, ['behavior:choice', 'behavior:probabilities', 'neural:A'])
+    expected = [Net().forward(_prompt_text(s))['early'] for s in ('x', 'yy', 'zzz')]
+    np.testing.assert_array_equal(session.collect('neural:A').values, expected)
+    assert session.collect('behavior:probabilities').attrs['measure']['source'] == 'likelihood'
+
+
+def test_a_composite_region_records_its_chosen_units():
+    net = Net()
+    region = CompositeSelector(layers=(('early', (0, 2)), ('late', None)), modality='text')
+    task = TaskContext(instruction='Which?', response=Choice(['cat', 'dog']))
+    recorded = _run(_model(net, {'R': region}), task, ['behavior:choice', 'neural:R']).collect('neural:R')
+    first = Net().forward(_prompt_text('x'))
+    np.testing.assert_array_equal(recorded.values[0], [first['early'][0], first['early'][2], *first['late']])
+
+
+def test_a_streaming_session_gets_each_trial_as_it_ends():
+    net = Net()
+    task = TaskContext(instruction='Which?', response=Choice(['cat', 'dog']))
+    session = _run(_model(net), task, ['behavior:choice', 'neural:A'], streaming=True)
+    assert [e.t_ms for e in session.emitted if e.channel == 'neural:A'] == [0.0, 1.0, 2.0]
+    assert session.collect('neural:A').sizes['presentation'] == 3
+
+
+def test_an_extractor_without_record_is_recorded_through_the_model():
+    torch = pytest.importorskip('torch')
+    network = torch.nn.Sequential(torch.nn.Linear(2, 3), torch.nn.ReLU(), torch.nn.Linear(3, 1))
+
+    def extract(stimuli, layers, **kwargs):
+        raise AssertionError('recorded during the task')
+
+    def reply(prompt):
+        network(torch.tensor([[float(len(prompt.text())), 1.0]]))
+        return 'cat'
+    model = BrainScoreModel('plain', model=network, extractor={'text': extract},
+                            region_layer_map={'A': LayerSelector('0', modality='text')}, reply=reply)
+    task = TaskContext(instruction='Which?', response=Choice(['cat', 'dog']))
+    recorded = _run(model, task, ['behavior:choice', 'neural:A']).collect('neural:A')
+    expected = network[0](torch.tensor([[float(len(_prompt_text('x'))), 1.0]])).detach().numpy()[0]
+    np.testing.assert_allclose(recorded.values[0], expected, rtol=1e-6)
+    assert not network[0]._forward_hooks

From 6ff9061418a5f3a0e768fdc94090e723e5d9a4dd Mon Sep 17 00:00:00 2001
From: Kartik Pradeepan <kartikspradeepan@gmail.com>
Date: Wed, 4 Nov 2026 12:41:05 -0500
Subject: [PATCH 3/6] feat(contract): declare subject traits and keep memory
 across a session

Subject gains traits (empty by default). A benchmark that needs one
lists it in required_traits, and pre-flight marks a subject without it
N/A. BrainScoreModel(memory='session') declares memory:session and
passes each earlier trial's prompt and reply to reply/likelihood as
prompt.history; reset() and the start of a new session clear it.
---
 brainscore_core/_evaluation/compatibility.py |  8 ++
 brainscore_core/_model/behavioral.py         | 25 ++++++-
 brainscore_core/_sessions/drivers.py         |  1 +
 brainscore_core/brainscore_model.py          | 14 ++++
 brainscore_core/contract.py                  | 10 +++
 brainscore_core/responses.py                 |  1 +
 tests/test_session_memory.py                 | 79 ++++++++++++++++++++
 7 files changed, 135 insertions(+), 3 deletions(-)
 create mode 100644 tests/test_session_memory.py

diff --git a/brainscore_core/_evaluation/compatibility.py b/brainscore_core/_evaluation/compatibility.py
index 222c52eb..6ce24ffb 100644
--- a/brainscore_core/_evaluation/compatibility.py
+++ b/brainscore_core/_evaluation/compatibility.py
@@ -160,6 +160,14 @@ def check_channel_compatibility(subject: Subject, benchmark) -> CompatibilityMat
             f"'{benchmark_id}' sends {_format_channels(bench_inputs)}."
         )
 
+    required_traits = set(getattr(benchmark, 'required_traits', None) or ())
+    missing_traits = required_traits - set(subject.traits) if required_traits else set()
+    if missing_traits:
+        raise CompatibilityError(
+            f"Benchmark '{benchmark_id}' requires subject traits {sorted(missing_traits)}; "
+            f"subject '{subject_id}' declares {sorted(subject.traits) or 'none'}."
+        )
+
     if getattr(benchmark, "uses_session", False):
         probe = getattr(subject, "check_session_support", None)
         if callable(probe):
diff --git a/brainscore_core/_model/behavioral.py b/brainscore_core/_model/behavioral.py
index 1ae95d47..413f18f4 100644
--- a/brainscore_core/_model/behavioral.py
+++ b/brainscore_core/_model/behavioral.py
@@ -1,6 +1,7 @@
 """Behavioral readout and generation support for BrainScoreModel."""
 
 from contextlib import ExitStack, contextmanager, nullcontext
+from dataclasses import replace
 from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple
 
 from ..contract import TaskContext
@@ -25,6 +26,8 @@ def __init__(self, owner) -> None:
         self.configured_from: Optional[TaskContext] = None
         self._configured_as: Optional[set] = None  # fingerprints of the fitted configuration
         self._stale = False
+        # With memory='session': the session's earlier trials as (prompt, reply).
+        self.history: list = []
         # Response tasks (TaskContext.response): how the answer and probabilities are produced.
         self.response_path: Optional[str] = None
         self.probability_source: Optional[str] = None
@@ -181,6 +184,7 @@ def start_task(self, task_context_or_task, fitting_stimuli=None,
             )
 
     def reset(self) -> None:
+        self.history = []
         self.response_path = self.probability_source = None
         self.task_context = None
         self.readout_classifier = None
@@ -343,14 +347,29 @@ def respond(self, stimuli, channels, regions=()):
                     outputs['behavior:probabilities'] = _implied(values, 'readout')
         else:
             score_first = path == 'generation' and self.probability_source == 'likelihood'
+            remember = owner._memory == 'session'
+            if remember and owner._reply_fn is None and owner._likelihood_fn is None:
+                raise ValueError("memory='session' needs reply= or likelihood=; answer= cannot see "
+                                 'earlier trials.')
+
+            def with_history(prompt):
+                return replace(prompt, history=tuple(self.history)) if remember else prompt
+
             if path == 'likelihood':
                 def call(row, prompt):
-                    return owner._likelihood_fn(prompt, list(labels))
+                    scores = owner._likelihood_fn(with_history(prompt), list(labels))
+                    if remember:
+                        self.history.append((prompt, max(scores, key=scores.get)))
+                    return scores
             elif owner._reply_fn is not None:
                 def call(row, prompt):
+                    shown = with_history(prompt)
                     # Scoring first keeps the recorded first pass the prompt alone.
-                    scores = owner._likelihood_fn(prompt, list(labels)) if score_first else None
-                    return owner._reply_fn(prompt), scores
+                    scores = owner._likelihood_fn(shown, list(labels)) if score_first else None
+                    reply = owner._reply_fn(shown)
+                    if remember:
+                        self.history.append((prompt, str(reply)))
+                    return reply, scores
             else:   # a registration's answer: it phrases and parses a Choice itself
                 def call(row, prompt):
                     return owner._generation_fn(stimulus_row=row, instruction=task.instruction,
diff --git a/brainscore_core/_sessions/drivers.py b/brainscore_core/_sessions/drivers.py
index 78f2f574..16a2ec96 100644
--- a/brainscore_core/_sessions/drivers.py
+++ b/brainscore_core/_sessions/drivers.py
@@ -403,6 +403,7 @@ def drive_behavior_and_neural_session(model, session, driver: str = "interact")
     answers = [c for c in requested if parse_channel(c)[0] == "behavior"]
     _refuse_variance(model, session)
     model.start_task(task_context)
+    model._behavioral.history = []     # memory spans one session, not the next
     stimuli = _behavior_stimuli_from_session(session, task_context)
     if stimuli is None:
         raise ValueError("Recording during a task needs stimuli to present.")
diff --git a/brainscore_core/brainscore_model.py b/brainscore_core/brainscore_model.py
index 7b43fb6b..87779dbf 100644
--- a/brainscore_core/brainscore_model.py
+++ b/brainscore_core/brainscore_model.py
@@ -96,6 +96,7 @@ def __init__(
         intervention: Optional[Callable] = None,
         reply: Optional[Callable] = None,
         likelihood: Optional[Callable] = None,
+        memory: Optional[str] = None,
         behavioral_readout_layer: Optional[str] = None,
         generation_fn: Optional[Callable] = None,
         action_fn: Optional[Callable] = None,
@@ -141,6 +142,12 @@ def __init__(
             probability}``; scores candidate replies without generating. Gives
             choices and their implied probabilities. Run the prompt on its own
             before any candidate, so a recording during scoring is the prompt pass.
+        :param memory: ``'session'`` keeps what the model saw and said across the
+            trials of a session, until ``reset()``: each trial's prompt and reply
+            reach ``reply`` and ``likelihood`` as ``prompt.history``. It declares
+            the ``memory:session`` trait. A stateful extractor (a recurrent model)
+            keeps its own state the same way. ``None`` (default): every trial is
+            independent.
         :param policy: ``EnvironmentStep -> EnvironmentResponse``; chooses an
             action for closed-loop / embodied evaluation.
         :param intervention: ``StateChange -> (PerturbationApplied, cleanup)``;
@@ -235,6 +242,9 @@ def __init__(
             'behavioral_readout_layer')
         self._generation_fn = capability_values.get('generation_fn')
         self._reply_fn, self._likelihood_fn = reply, likelihood
+        if memory not in (None, 'session'):
+            raise ValueError("memory must be None or 'session'")
+        self._memory = memory
         self._action_fn = capability_values.get('action_fn')
         self._state_change_fn = capability_values.get('state_change_fn')
 
@@ -290,6 +300,10 @@ def _visual_forms(self):
         forms = {'vision': 'vision:still', 'video': 'vision:clip'}
         return {forms[key] for key in self._preprocessors if key in forms}
 
+    @property
+    def traits(self):
+        return {'memory:session'} if self._memory == 'session' else set()
+
     @property
     def out_channels(self):
         # Every output comes from a registered part; nothing is found by name.
diff --git a/brainscore_core/contract.py b/brainscore_core/contract.py
index ae632297..ff3fbd91 100644
--- a/brainscore_core/contract.py
+++ b/brainscore_core/contract.py
@@ -127,6 +127,16 @@ def required_channels(self) -> Set[str]:
         """Inputs the subject cannot run without; none by default."""
         return set()
 
+    @property
+    def traits(self) -> Set[str]:
+        """Properties of the subject a benchmark may require; none by default.
+
+        ``'memory:session'``: the subject keeps what it saw and said across the
+        trials of one session, until ``reset()``. A benchmark that needs a trait
+        lists it in ``required_traits``; a subject without it is N/A.
+        """
+        return set()
+
     @abstractmethod
     def participate(self, session: "Session") -> None:
         """Consume session inputs and emit responses until the session ends."""
diff --git a/brainscore_core/responses.py b/brainscore_core/responses.py
index 294aaa27..c7885579 100644
--- a/brainscore_core/responses.py
+++ b/brainscore_core/responses.py
@@ -127,6 +127,7 @@ class Prompt:
     parts: Tuple[Tuple[str, Any], ...]
     stimulus_id: Optional[str] = None
     names: Tuple[str, ...] = ()   # one per part: 'instruction', a column name, or 'directive'
+    history: Tuple[Tuple['Prompt', str], ...] = ()   # earlier trials of the session: (prompt, reply)
 
     def text(self) -> str:
         """The text parts joined, for a model that reads text only."""
diff --git a/tests/test_session_memory.py b/tests/test_session_memory.py
new file mode 100644
index 00000000..9497a113
--- /dev/null
+++ b/tests/test_session_memory.py
@@ -0,0 +1,79 @@
+"""Memory across trials: a declared trait, required by benchmarks, kept through a session."""
+import pytest
+
+from brainscore_core.brainscore_model import BrainScoreModel
+from brainscore_core.compatibility import CompatibilityError
+from brainscore_core.model_interface import TaskContext
+from brainscore_core.responses import Choice
+from brainscore_core.streaming import BufferedSession
+from brainscore_core.supported_data_standards.brainio.stimuli import StimulusSet
+from brainscore_core._evaluation.compatibility import check_channel_compatibility
+
+pytestmark = pytest.mark.unit
+
+WORDS = ['cat', 'cat', 'dog', 'sun', 'sun', 'dog']
+ONE_BACK = [None, 'yes', 'no', 'no', 'yes', 'no']        # same as the previous trial?
+STIMULI = StimulusSet({'stimulus_id': [f's{i}' for i in range(6)], 'sentence': WORDS})
+STIMULI.identifier = 'one-back'
+TASK = TaskContext(instruction='Is this word the same as the previous one?',
+                   response=Choice(['yes', 'no']))
+
+
+def _word(prompt):
+    return [content for kind, content in prompt.parts if kind == 'text'][1]
+
+
+def one_back(prompt):
+    """Answers from the session's history; with no history it can only say no."""
+    previous = _word(prompt.history[-1][0]) if prompt.history else None
+    return 'yes' if previous == _word(prompt) else 'no'
+
+
+def _answers(model):
+    session = BufferedSession(stimulus_set=STIMULI, task_context=TASK,
+                              requested_output_channels=['behavior:choice'])
+    model.participate(session)
+    return list(session.collect('behavior:choice')['raw_answer'].values)
+
+
+def test_a_model_with_memory_sees_earlier_trials():
+    model = BrainScoreModel('remembers', supported_modalities={'text'}, reply=one_back, memory='session')
+    assert model.traits == {'memory:session'}
+    assert _answers(model)[1:] == ONE_BACK[1:]
+
+
+def test_without_memory_every_trial_is_alone():
+    model = BrainScoreModel('forgets', supported_modalities={'text'}, reply=one_back)
+    assert model.traits == set()
+    assert _answers(model) == ['no'] * 6
+
+
+def test_memory_lasts_one_session():
+    model = BrainScoreModel('remembers', supported_modalities={'text'}, reply=one_back, memory='session')
+    _answers(model)
+    assert len(model._behavioral.history) == 6
+    model.reset()
+    assert model._behavioral.history == []
+    assert _answers(model)[1:] == ONE_BACK[1:]               # a new session starts fresh
+
+
+class _MemoryBenchmark:
+    identifier = 'one-back'
+    required_input_channels = {'text'}
+    requested_output_channels = {'behavior:choice'}
+    required_traits = {'memory:session'}
+
+
+def test_a_memory_benchmark_marks_a_model_without_memory_na():
+    forgets = BrainScoreModel('forgets', supported_modalities={'text'}, reply=one_back)
+    with pytest.raises(CompatibilityError, match='memory:session'):
+        check_channel_compatibility(forgets, _MemoryBenchmark())
+    remembers = BrainScoreModel('remembers', supported_modalities={'text'}, reply=one_back, memory='session')
+    check_channel_compatibility(remembers, _MemoryBenchmark())
+
+
+def test_memory_needs_a_part_that_can_see_history():
+    model = BrainScoreModel('answers', supported_modalities={'text'}, memory='session',
+                            answer=lambda stimulus_row, instruction, label_set: 'no')
+    with pytest.raises(ValueError, match='reply'):
+        _answers(model)

From 84d05908be01276854cb57919b524e09161b0124 Mon Sep 17 00:00:00 2001
From: Kartik Pradeepan <kartikspradeepan@gmail.com>
Date: Wed, 4 Nov 2026 12:54:58 -0500
Subject: [PATCH 4/6] feat(model): accept fused extractors for several input
 types at once

extractor={('vision', 'text'): fused} declares each type and is called
once for a stimulus that brings several of them; single-type
extractors still take single-type stimuli.
---
 brainscore_core/_model/behavioral.py   |  2 +
 brainscore_core/_model/dispatch.py     | 30 ++++++++++++-
 brainscore_core/brainscore_model.py    |  5 +++
 brainscore_core/capabilities/neural.py |  3 +-
 tests/test_fused_extractor.py          | 62 ++++++++++++++++++++++++++
 5 files changed, 100 insertions(+), 2 deletions(-)
 create mode 100644 tests/test_fused_extractor.py

diff --git a/brainscore_core/_model/behavioral.py b/brainscore_core/_model/behavioral.py
index 413f18f4..2d9134e4 100644
--- a/brainscore_core/_model/behavioral.py
+++ b/brainscore_core/_model/behavioral.py
@@ -530,6 +530,8 @@ def _region_modality(self, region: str) -> str:
         if modality is not None:
             return canonical_modality(modality)
         modalities = {canonical_modality(key) for key in self.owner._extractors}
+        if len({id(e) for e in self.owner._extractors.values()}) == 1:
+            return sorted(modalities)[0]        # one (fused) extractor records every region
         if len(modalities) != 1:
             raise ValueError(
                 f'Region {region!r} does not say which input it reads. Map it with '
diff --git a/brainscore_core/_model/dispatch.py b/brainscore_core/_model/dispatch.py
index 8895c4b3..416e7c78 100644
--- a/brainscore_core/_model/dispatch.py
+++ b/brainscore_core/_model/dispatch.py
@@ -69,6 +69,16 @@ def _is_clip(stimuli) -> bool:
     return any(column in getattr(stimuli, 'columns', ()) for column in CLIP_COLUMNS)
 
 
+def fused_for(owner, detected):
+    """The fused extractor taking every detected input type, if one is registered."""
+    if len(detected) < 2:
+        return None
+    for types, extractor in owner._fused_extractors:
+        if set(detected) <= types:
+            return extractor
+    return None
+
+
 def extractor_for(owner, modality: str, stimuli):
     """The extractor for ``modality``; with a still and a clip extractor, the one
     matching the stimulus's form. None when the model has no such extractor."""
@@ -93,7 +103,18 @@ def normalize_extractors(owner, preprocessors, activations_model):
     import inspect
 
     if isinstance(activations_model, Mapping):
-        extractors = dict(activations_model)
+        # A fused extractor is registered under every input type it takes at
+        # once, {('vision', 'text'): fused}. It also serves each type alone
+        # unless a single-type extractor is registered for it.
+        fused = [(frozenset(canonical_modality(m) for m in key), value)
+                 for key, value in activations_model.items() if isinstance(key, tuple)]
+        extractors = {}
+        for types, value in fused:
+            for modality in types:
+                extractors.setdefault(modality, value)
+        extractors.update({key: value for key, value in activations_model.items()
+                           if not isinstance(key, tuple)})
+        owner._fused_extractors = fused
     elif activations_model is not None:
         if not preprocessors:
             raise ValueError(
@@ -283,6 +304,13 @@ def process_multi_modality(self, stimuli, detected: Set[str],
         import xarray as xr
 
         owner = self.owner
+        fused = fused_for(owner, detected)
+        if fused is not None:
+            # One fused extractor takes every input type: one forward pass.
+            assembly = fused(stimuli, layers=layers)
+            if owner._recorder.is_multi_region:
+                assembly = owner._tag_neuroids_with_regions(assembly)
+            return assembly
         ordered = sorted(
             detected,
             key=lambda m: (
diff --git a/brainscore_core/brainscore_model.py b/brainscore_core/brainscore_model.py
index 87779dbf..fe93b875 100644
--- a/brainscore_core/brainscore_model.py
+++ b/brainscore_core/brainscore_model.py
@@ -133,6 +133,10 @@ def __init__(
         :param answer: ``(stimulus_row, instruction, label_set) -> label``;
             generates and parses a behavioral answer for instruction-following
             tasks.
+        :param extractor: also accepts a fused extractor under a tuple of input
+            types, ``{('vision', 'text'): fused}``: one forward pass over all of a
+            stimulus's parts. It declares each type, and takes any stimulus that
+            brings several of them at once.
         :param reply: ``(prompt) -> str``; the model's raw reply to one trial.
             ``prompt`` is a ``brainscore_core.responses.Prompt``: the trial's parts
             in order (instruction, stimulus parts, the response's directive).
@@ -228,6 +232,7 @@ def __init__(
         from ._model.dispatch import normalize_extractors
         self._preprocessors: Dict[str, Optional[Callable]] = dict(preprocessors)
         self._activations_model = activations_model
+        self._fused_extractors = []     # (input types, extractor) registered under a tuple
         self._extractors, self._layer_extractors = normalize_extractors(
             self, preprocessors, activations_model)
         for modality in supported_modalities or ():
diff --git a/brainscore_core/capabilities/neural.py b/brainscore_core/capabilities/neural.py
index 423f3b95..6a8dc7fc 100644
--- a/brainscore_core/capabilities/neural.py
+++ b/brainscore_core/capabilities/neural.py
@@ -65,7 +65,8 @@ def process(self, model, stimuli, **kwargs):
                 f"{sorted(model._region_layer_map_dict.keys())}."
             )
 
-        if multi_modality and len(detected) > 1:
+        from .._model.dispatch import fused_for
+        if (multi_modality or fused_for(model, detected) is not None) and len(detected) > 1:
             return model._process_multi_modality(stimuli, detected, layers)
 
         if not multi_modality and len(detected) > 1:
diff --git a/tests/test_fused_extractor.py b/tests/test_fused_extractor.py
new file mode 100644
index 00000000..18105ed6
--- /dev/null
+++ b/tests/test_fused_extractor.py
@@ -0,0 +1,62 @@
+"""A fused extractor: one forward pass over several input types, split by input."""
+import numpy as np
+import pytest
+
+from brainscore_core.brainscore_model import BrainScoreModel
+from brainscore_core.selection import LayerSelector
+from brainscore_core.streaming import BufferedSession
+from brainscore_core.supported_data_standards.brainio.assemblies import NeuroidAssembly
+from brainscore_core.supported_data_standards.brainio.stimuli import StimulusSet
+
+pytestmark = pytest.mark.unit
+
+
+class Fused:
+    """Sees the image and the sentence of each stimulus together; returns one step per input."""
+    identifier = backbone_id = 'fused'
+
+    def __init__(self):
+        self.calls = []
+
+    def __call__(self, stimuli, layers, **kwargs):
+        self.calls.append(sorted(c for c in stimuli.columns if c != 'stimulus_id'))
+        ids = list(stimuli['stimulus_id'])
+        values = np.stack([[[1.0, 2.0], [3.0, 4.0]] for _ in ids])       # parts: image, text
+        return NeuroidAssembly(values, dims=['presentation', 'part', 'neuroid'], coords={
+            'stimulus_id': ('presentation', ids), 'part_name': ('part', ['image', 'text']),
+            'part_index': ('part', [0, 1]),
+            'neuroid_id': ('neuroid', ['n0', 'n1']), 'layer': ('neuroid', [layers[0]] * 2)})
+
+
+def test_a_fused_extractor_declares_each_input_and_runs_once():
+    fused = Fused()
+    model = BrainScoreModel('fused', extractor={('vision', 'text'): fused},
+                            region_layer_map={'A': 'layers.3'})
+    assert {'vision', 'text', 'vision:still'} <= model.in_channels
+    both = StimulusSet({'stimulus_id': ['a', 'b'], 'image_path': ['a.png', 'b.png'],
+                        'sentence': ['one', 'two']})
+    both.identifier = 'fused-both'
+    session = BufferedSession(stimulus_set=both, requested_output_channels=['neural:A'])
+    model.participate(session)
+    recorded = session.collect('neural:A')
+    assert len(fused.calls) == 1 and fused.calls[0] == ['image_path', 'sentence']
+    assert recorded.dims == ('presentation', 'part', 'neuroid')
+    assert list(recorded['part_name'].values) == ['image', 'text']
+
+
+def test_single_input_extractors_still_take_single_input_stimuli():
+    fused, vision_calls = Fused(), []
+
+    def vision(stimuli, layers, **kwargs):
+        vision_calls.append(1)
+        ids = list(stimuli['stimulus_id'])
+        return NeuroidAssembly(np.ones((len(ids), 2)), dims=['presentation', 'neuroid'], coords={
+            'stimulus_id': ('presentation', ids), 'neuroid_id': ('neuroid', ['v0', 'v1']),
+            'layer': ('neuroid', [layers[0]] * 2)})
+    model = BrainScoreModel('mixed', extractor={'vision': vision, ('vision', 'text'): fused},
+                            region_layer_map={'A': LayerSelector('layers.3', modality='vision')})
+    images = StimulusSet({'stimulus_id': ['a'], 'image_path': ['a.png']})
+    images.identifier = 'images-only'
+    session = BufferedSession(stimulus_set=images, requested_output_channels=['neural:A'])
+    model.participate(session)
+    assert session.collect('neural:A').dims == ('presentation', 'neuroid')

From af0ca920b132d1f563a1036fe39452439fbaccd1 Mon Sep 17 00:00:00 2001
From: Kartik Pradeepan <kartikspradeepan@gmail.com>
Date: Wed, 4 Nov 2026 13:04:01 -0500
Subject: [PATCH 5/6] feat(tasks): give choices implied reaction times

Choice(labels, reaction_time='evidence_margin' | 'surprisal') returns
an implied RT per stimulus on behavior:rt, from the readout's or the
likelihood's label probabilities, labelled with its link. Readout and
likelihood parts declare behavior:rt.
---
 brainscore_core/_model/behavioral.py     | 19 ++++++++++++++++++-
 brainscore_core/capabilities/sessions.py |  4 ++--
 brainscore_core/io_catalog.py            |  2 +-
 brainscore_core/responses.py             | 13 ++++++++++++-
 tests/test_compatibility.py              |  2 +-
 tests/test_subject_channels.py           |  2 +-
 tests/test_task_responses.py             | 17 +++++++++++++++++
 7 files changed, 52 insertions(+), 7 deletions(-)

diff --git a/brainscore_core/_model/behavioral.py b/brainscore_core/_model/behavioral.py
index 2d9134e4..7ad05e1a 100644
--- a/brainscore_core/_model/behavioral.py
+++ b/brainscore_core/_model/behavioral.py
@@ -302,7 +302,7 @@ def _start_response_task(self, task_context) -> None:
                 + (' with reported confidence (needs reply=)' if reported_confidence else '')
                 + '. Paths tried: ' + '; '.join(f'{p} needs {needs[p]}' for p in order))
         source = None
-        if isinstance(response, Choice) and response.probabilities:
+        if isinstance(response, Choice) and (response.probabilities or response.reaction_time):
             if path != 'readout' and owner._likelihood_fn is not None:
                 source = 'likelihood'
             elif readout_wired:
@@ -410,6 +410,9 @@ def call(row, prompt):
                 elif self.probability_source == 'readout':
                     outputs['behavior:probabilities'] = _implied(self.readout_classifier.predict_proba(
                         owner._extract_behavioral_features(stimuli)), 'readout')
+        if isinstance(response, Choice) and response.reaction_time:
+            outputs['behavior:rt'] = _reaction_time(outputs['behavior:probabilities'],
+                                                    response.reaction_time, stimuli)
         outputs['behavior'] = outputs[response.channels[0]]
         if prompts:
             # A multi-image trial records one step per image: name each by its part.
@@ -759,6 +762,20 @@ def _texts(parsed, replies, stimuli):
                               attrs={'measure': {'kind': 'reported', 'source': 'reply'}})
 
 
+def _reaction_time(probabilities, link, stimuli):
+    """An implied reaction time per stimulus from label probabilities (larger = slower)."""
+    from brainscore_core.supported_data_standards.brainio.assemblies import BehavioralAssembly
+    values = np.asarray(probabilities.values, dtype=float)
+    ordered = np.sort(values, axis=1)[:, ::-1]
+    if link == 'evidence_margin':
+        rt = 1.0 - (ordered[:, 0] - ordered[:, 1])
+    else:                                         # surprisal of the chosen (most probable) label
+        rt = -np.log2(np.clip(ordered[:, 0], 1e-12, 1.0))
+    return BehavioralAssembly(rt, dims=['presentation'], coords=_presentation(stimuli),
+                              attrs={'measure': {'kind': 'implied', 'link': link,
+                                                 'source': probabilities.attrs.get('measure', {}).get('source')}})
+
+
 def _softmax_rows(scores, labels):
     logp = np.array([[float(row[label]) for label in labels] for row in scores], dtype=float)
     logp -= logp.max(axis=1, keepdims=True)
diff --git a/brainscore_core/capabilities/sessions.py b/brainscore_core/capabilities/sessions.py
index a0ae6265..97c0b855 100644
--- a/brainscore_core/capabilities/sessions.py
+++ b/brainscore_core/capabilities/sessions.py
@@ -71,14 +71,14 @@ def declared_output_channels(self, model):
         # Each part declares the kinds of answer it can give (rule 2.4).
         channels = {'behavior'}
         if model._behavioral_readout_layer is not None:
-            channels |= {'behavior:choice', 'behavior:probabilities', 'behavior:rating'}
+            channels |= {'behavior:choice', 'behavior:probabilities', 'behavior:rating', 'behavior:rt'}
         if model._generation_fn is not None:
             channels |= {'behavior:choice'}
         if model._reply_fn is not None:
             channels |= {'behavior:choice', 'behavior:rating', 'behavior:confidence',
                          'behavior:text'}
         if model._likelihood_fn is not None:
-            channels |= {'behavior:choice', 'behavior:probabilities'}
+            channels |= {'behavior:choice', 'behavior:probabilities', 'behavior:rt'}
         return channels
 
     def supports_session(self, model, channels) -> bool:
diff --git a/brainscore_core/io_catalog.py b/brainscore_core/io_catalog.py
index e758f122..a56591a4 100644
--- a/brainscore_core/io_catalog.py
+++ b/brainscore_core/io_catalog.py
@@ -363,7 +363,7 @@ def validate(
                  "a fitted readout or the model's generation path",
                  meta_keys=("label_set",),
                  addressing=("address is the kind of answer a TaskContext response asks for: "
-                             "choice, probabilities, confidence, rating or text"),
+                             "choice, probabilities, confidence, rating, rt or text"),
                  materializers=("start_task", "BehavioralAssembly")),
     # --- outputs: motor (carried by EnvironmentResponse, harness-defined) ---
     CatalogEntry("motor", OUTPUT, "EnvironmentResponse.action",
diff --git a/brainscore_core/responses.py b/brainscore_core/responses.py
index c7885579..8b9433d7 100644
--- a/brainscore_core/responses.py
+++ b/brainscore_core/responses.py
@@ -51,21 +51,32 @@ class Choice:
     ``probabilities=True`` also returns the implied probability of each label
     (from a readout or a language model's likelihoods). ``confidence=Scale(...)``
     also asks the subject to report how sure it is, in the same reply.
+    ``reaction_time`` also returns an implied reaction time from those
+    probabilities, scored by correlation with human RTs as Futrell2018 scores
+    surprisal: ``'evidence_margin'`` (1 minus the gap between the two most
+    probable labels; close calls are slow) or ``'surprisal'`` (-log2 of the
+    chosen label's probability).
     """
     labels: Tuple[str, ...]
     probabilities: bool = False
     confidence: Optional[Scale] = None
+    reaction_time: Optional[str] = None   # 'evidence_margin' or 'surprisal': an implied RT
+
+    REACTION_TIMES = ('evidence_margin', 'surprisal')
 
     def __post_init__(self):
         if isinstance(self.labels, str) or len(self.labels) < 2:
             raise ValueError('Choice needs at least two labels')
+        if self.reaction_time not in (None, *self.REACTION_TIMES):
+            raise ValueError(f'reaction_time is one of {self.REACTION_TIMES}, not {self.reaction_time!r}')
         object.__setattr__(self, 'labels', tuple(self.labels))
 
     @property
     def channels(self) -> Tuple[str, ...]:
         return (('behavior:choice',)
                 + (('behavior:probabilities',) if self.probabilities else ())
-                + (('behavior:confidence',) if self.confidence is not None else ()))
+                + (('behavior:confidence',) if self.confidence is not None else ())
+                + (('behavior:rt',) if self.reaction_time is not None else ()))
 
     def directive(self) -> str:
         text = f'Answer with exactly one of: {", ".join(self.labels)}.'
diff --git a/tests/test_compatibility.py b/tests/test_compatibility.py
index e586bd77..572c07c3 100644
--- a/tests/test_compatibility.py
+++ b/tests/test_compatibility.py
@@ -395,7 +395,7 @@ def _model(self):
     def test_every_part_declares_registered_channels(self):
         model = self._model()
         assert model.out_channels == {'neural:IT', 'behavior', 'behavior:choice',
-                                      'behavior:probabilities', 'behavior:rating',
+                                      'behavior:probabilities', 'behavior:rating', 'behavior:rt',
                                       'motor', 'perturbation'}
         assert {'vision', 'observation', 'lesion', 'stimulation'} <= model.in_channels
         for output in sorted(model.out_channels):
diff --git a/tests/test_subject_channels.py b/tests/test_subject_channels.py
index a7dd7ff0..bba6b0d2 100644
--- a/tests/test_subject_channels.py
+++ b/tests/test_subject_channels.py
@@ -100,7 +100,7 @@ def test_brainscore_model_gets_channel_identity_from_preprocessors():
 
     assert model.in_channels == {"vision", "vision:still"}   # a vision key takes stills
     assert model.out_channels == {"neural:IT", "behavior", "behavior:choice",
-                                  "behavior:probabilities", "behavior:rating"}
+                                  "behavior:probabilities", "behavior:rating", "behavior:rt"}
     assert model.required_channels == set()
 
 
diff --git a/tests/test_task_responses.py b/tests/test_task_responses.py
index 663af060..6db13331 100644
--- a/tests/test_task_responses.py
+++ b/tests/test_task_responses.py
@@ -181,3 +181,20 @@ def test_named_parts_show_a_sequence_in_order():
         trial_prompt(TaskContext(response=Choice(['yes', 'no']), present=('image:word_3',)), row)
     with pytest.raises(ValueError, match='kind'):
         TaskContext(response=Choice(['yes', 'no']), present=('picture:word_1',))
+
+
+@pytest.mark.parametrize('link', ['evidence_margin', 'surprisal'])
+def test_implied_reaction_times_rise_for_close_calls(link):
+    scores = {'a': {'cat': -0.1, 'dog': -5.0}, 'b': {'cat': -1.0, 'dog': -1.2},
+              'c': {'cat': -3.0, 'dog': -0.05}, 'd': {'cat': -0.7, 'dog': -0.69}}
+    model, _ = _model(scores=scores)
+    task = TaskContext(instruction='Which?', response=Choice(['cat', 'dog'], reaction_time=link))
+    rt = _ask(model, task, ['behavior:rt']).collect('behavior:rt')
+    assert rt.attrs['measure'] == {'kind': 'implied', 'link': link, 'source': 'likelihood'}
+    fast, slow = rt.values[0], rt.values[3]                   # a clear call, a near tie
+    assert slow > fast
+
+
+def test_a_reply_only_model_cannot_give_reaction_times():
+    model, _ = _model(replies={s: 'cat' for s in 'abcd'})
+    assert 'behavior:rt' not in model.out_channels

From 69d01b49a5652fee6019eef3f9f9ca0911e7e435 Mon Sep 17 00:00:00 2001
From: Kartik Pradeepan <kartikspradeepan@gmail.com>
Date: Wed, 4 Nov 2026 13:16:20 -0500
Subject: [PATCH 6/6] ci: run unified tests from the trial sequences branch

---
 .github/workflows/umi-v2.yml | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/.github/workflows/umi-v2.yml b/.github/workflows/umi-v2.yml
index 438329b9..6bdfe354 100644
--- a/.github/workflows/umi-v2.yml
+++ b/.github/workflows/umi-v2.yml
@@ -11,9 +11,9 @@ jobs:
   coordinated:
     name: UMI v2
     # Update the workflow and runner revision together.
-    uses: KartikP/brainscore-unified/.github/workflows/offline-tests.yml@fd3bec0605731b9c72900dfd7e85137a3aa56312
+    uses: KartikP/brainscore-unified/.github/workflows/offline-tests.yml@346b90ab5093d520d4c49c10752acce0afa6fb39
     with:
       candidate: core
       candidate_repository: ${{ github.event.pull_request.head.repo.full_name || github.repository }}
       candidate_revision: ${{ github.event.pull_request.head.sha || github.sha }}
-      unified_revision: fd3bec0605731b9c72900dfd7e85137a3aa56312
+      unified_revision: 346b90ab5093d520d4c49c10752acce0afa6fb39