code-loader 1.0.193.dev4__py3-none-any.whl → 1.0.193.dev5__py3-none-any.whl

code_loader/leaploader.py

@@ -7,7 +7,7 @@ import sys
 from contextlib import redirect_stdout
 from functools import lru_cache
 from pathlib import Path
-from typing import Dict, List, Iterable, Set, FrozenSet, Union, Any, Type, Optional, Callable, Tuple
+from typing import Dict, List, Iterable, Set, Union, Any, Type, Optional, Callable, Tuple
 
 import numpy as np
 import numpy.typing as npt
@@ -27,8 +27,7 @@ from code_loader.contract.sim_config import FloatBounds, IntBounds, CategoricalB
 from code_loader.inner_leap_binder import global_leap_binder
 from code_loader.inner_leap_binder.leapbinder import mapping_runtime_mode_env_var_mame
 from code_loader.leaploaderbase import LeapLoaderBase
-from code_loader.utils import get_root_exception_file_and_line_number, get_metadata_type_from_variable, \
-    validate_autoregressive_state_types, autoregressive_nests_equal
+from code_loader.utils import get_root_exception_file_and_line_number, get_metadata_type_from_variable
 
 
 def _serialize_sim_bounds(bounds) -> dict:
@@ -60,14 +59,35 @@ class LeapLoader(LeapLoaderBase):
 
     @lru_cache()
     def exec_script(self) -> None:
+        from code_loader.inner_leap_binder import leapbinder_decorators as _leap_dec
         try:
             os.environ[mapping_runtime_mode_env_var_mame] = 'TRUE'
             self.evaluate_module()
+
+            # A grouped integration test may index a grouped input/gt (image_list[i]) to wire a
+            # single sample of a group; that indexing is only valid when the mapping graph is
+            # built in grouped mode. Detect grouped-ness so the placeholder pass below matches
+            # the dataset. Preprocess is stubbed while mapping mode is on, so run it with mapping
+            # mode momentarily disabled, and cache it so _preprocess_result() doesn't re-run it.
+            is_grouped = False
             if global_leap_binder.integration_test_func is not None:
-                from code_loader.inner_leap_binder.leapbinder_decorators import \
-                    _reset_model_loop_state
-                _reset_model_loop_state()
-                global_leap_binder.integration_test_func(None, PreprocessResponse(state=DataStateType.training, length=0))
+                del os.environ[mapping_runtime_mode_env_var_mame]
+                try:
+                    preprocess_result = global_leap_binder.get_preprocess_result()
+                    self._preprocess_result_cached = preprocess_result
+                    is_grouped = any(r.is_grouped for r in preprocess_result.values())
+                except Exception:
+                    is_grouped = False
+                finally:
+                    os.environ[mapping_runtime_mode_env_var_mame] = 'TRUE'
+
+            _leap_dec._mapping_dataset_is_grouped = is_grouped
+            if global_leap_binder.integration_test_func is not None:
+                mapping_preprocess = (
+                    PreprocessResponse(sample_ids=[["__mapping_placeholder__"]], state=DataStateType.training)
+                    if is_grouped else
+                    PreprocessResponse(state=DataStateType.training, length=0))
+                global_leap_binder.integration_test_func(None, mapping_preprocess)
         except TypeError as e:
             import traceback
             global_leap_binder.setup_container = DatasetIntegrationSetup()
@@ -80,6 +100,7 @@ class LeapLoader(LeapLoaderBase):
             raise DatasetScriptException(getattr(e, 'message', repr(e))) from e
         finally:
             # ensure that the environment variable is removed after the script execution
+            _leap_dec._mapping_dataset_is_grouped = False
             if mapping_runtime_mode_env_var_mame in os.environ:
                 del os.environ[mapping_runtime_mode_env_var_mame]
 
@@ -250,16 +271,20 @@ class LeapLoader(LeapLoaderBase):
             )
 
         preprocess_result = self._preprocess_result()
-        if state == DataStateEnum.unlabeled and sample_id not in preprocess_result[state].sample_ids:
+        if state == DataStateEnum.unlabeled and sample_id not in preprocess_result[state].flat_sample_ids:
             self._preprocess_result(update_unlabeled_preprocess=True)
 
         metadata, metadata_is_none = self.get_metadata(state, sample_id)
 
         custom_latent_space = None
         if global_leap_binder.setup_container.custom_latent_space is not None:
-            custom_latent_space = global_leap_binder.setup_container.custom_latent_space.function(sample_id,
-                                                                                                  preprocess_result[
-                                                                                                      state])
+            latent_fn = global_leap_binder.setup_container.custom_latent_space.function
+            preprocess_state = preprocess_result[state]
+            if preprocess_state.is_grouped:
+                group_ids, pos = self._locate_group(preprocess_state, sample_id)
+                custom_latent_space = self._to_grouped_list(latent_fn(group_ids, preprocess_state))[pos]
+            else:
+                custom_latent_space = latent_fn(sample_id, preprocess_state)
         instance_mask = self._get_instances_masks(state, sample_id, instance_id)
         sample = DatasetSample(inputs=self._get_inputs(state, sample_id),
                                gt=None if state == DataStateEnum.unlabeled else self._get_gt(state, sample_id),
@@ -390,7 +415,7 @@ class LeapLoader(LeapLoaderBase):
             if self.get_sample_id_type() is str:
                 max_allowed_item_size = np.dtype('<U256').itemsize
                 for state, preprocess_response in preprocess_result.items():
-                    sample_ids_array = np.array(preprocess_response.sample_ids)
+                    sample_ids_array = np.array(preprocess_response.flat_sample_ids)
                     if sample_ids_array.dtype.itemsize > max_allowed_item_size:
                         raise Exception(f"Sample id are too long. Max allowed length is 256 charecters.")
 
@@ -480,49 +505,31 @@ class LeapLoader(LeapLoaderBase):
         try:
             preprocess_result = self._preprocess_result()
             input_shapes: Dict[str, List[int]] = {}
-            expected_keys: Optional[FrozenSet[str]] = None
-            expected_keys_state = ''
             for state, preprocess_response in preprocess_result.items():
                 if preprocess_response.sample_ids_to_instance_mappings:
                     raise Exception('Element instances are not supported together with '
                                     'tensorleap_autoregressive_step.')
                 sample_id = preprocess_response.sample_ids[0]
-                first_result = handler.function(sample_id, None, None, None, preprocess_response)
-                second_result = handler.function(sample_id, None, None, None, preprocess_response)
-                if not isinstance(first_result, tuple) or len(first_result) != 2:
-                    raise Exception('The autoregressive step hook must return a (next_inputs, '
-                                    f'state) tuple, got {type(first_result)} '
-                                    f'(state: {state.name}).')
-                first_inputs, first_state = first_result
-                if not isinstance(first_inputs, dict) or not first_inputs:
+                first_result = handler.function(sample_id, None, None, preprocess_response)
+                second_result = handler.function(sample_id, None, None, preprocess_response)
+                if not isinstance(first_result, dict):
                     raise Exception('The autoregressive step hook must return a dict of initial model '
-                                    f'inputs on its first call, got {type(first_inputs)} '
+                                    f'inputs on its first call, got {type(first_result)} '
                                     f'(state: {state.name}).')
-                validate_autoregressive_state_types(first_state)
-                if not isinstance(second_result, tuple) or len(second_result) != 2 or \
-                        not isinstance(second_result[0], dict) or \
-                        set(second_result[0].keys()) != set(first_inputs.keys()) or \
-                        not autoregressive_nests_equal(first_state, second_result[1]):
+                if not isinstance(second_result, dict) or \
+                        set(second_result.keys()) != set(first_result.keys()):
                     raise Exception('The autoregressive step hook is not deterministic — two calls '
                                     'with identical arguments returned different results '
                                     f'(state: {state.name}). Seed any stochasticity from '
                                     'sample_id: the engine replays steps after crash recovery '
                                     'and relies on identical results.')
-                if expected_keys is None:
-                    expected_keys = frozenset(first_inputs)
-                    expected_keys_state = state.name
-                elif frozenset(first_inputs) != expected_keys:
-                    raise Exception(f'The autoregressive step hook returned different model input '
-                                    f'keys across states: {sorted(expected_keys)} '
-                                    f'({expected_keys_state}) vs {sorted(first_inputs)} '
-                                    f'({state.name}). The key set must be identical for every '
-                                    f'state.')
-                for key, value in first_inputs.items():
-                    if key.startswith('_'):
-                        raise Exception(f'The autoregressive step hook returned model input key '
-                                        f'"{key}" — underscore passthrough keys were replaced by '
-                                        f'the explicit state channel; move passthrough values '
-                                        f'into the returned state. State: {state.name}.')
+                model_input_items = {key: value for key, value in first_result.items()
+                                     if not key.startswith('_')}
+                if not model_input_items:
+                    raise Exception('The autoregressive step hook returned no model inputs on its first '
+                                    'call — every key starts with "_" (underscore keys are passthrough '
+                                    f'state, never fed to the model). State: {state.name}.')
+                for key, value in model_input_items.items():
                     if not isinstance(value, np.ndarray):
                         raise Exception(f'The autoregressive step hook returned a non-numpy value for '
                                         f'model input "{key}" ({type(value)}). State: {state.name}.')
@@ -531,7 +538,7 @@ class LeapLoader(LeapLoaderBase):
                                         f'model input "{key}" across states: {input_shapes[key]} vs '
                                         f'{list(value.shape)}. Shapes must be fixed.')
                     input_shapes[key] = list(value.shape)
-                    if not autoregressive_nests_equal(value, second_result[0][key]):
+                    if not np.array_equal(value, second_result[key]):
                         raise Exception('The autoregressive step hook is not deterministic — two calls '
                                         'with identical arguments returned different outputs for model '
                                         f'input "{key}". Seed any stochasticity from sample_id: the '
@@ -545,39 +552,19 @@ class LeapLoader(LeapLoaderBase):
             test_result.is_passed = False
         return test_result
 
-    @staticmethod
-    def _validate_autoregressive_step_inputs(step_inputs: Dict[str, Any],
-                                             sample_id: Union[int, str]) -> None:
-        # The parse-time check only covers the first sample of each state; the engine runtime
-        # paths must enforce the same key/value contract for every sample they touch.
-        for key, value in step_inputs.items():
-            if key.startswith('_'):
-                raise Exception(f'The autoregressive step hook returned model input key "{key}" '
-                                f'for sample {sample_id} — underscore passthrough keys were '
-                                'replaced by the explicit state channel; move passthrough values '
-                                'into the returned state.')
-            if not isinstance(value, np.ndarray):
-                raise Exception('The autoregressive step hook returned a non-numpy value for '
-                                f'model input "{key}" ({type(value).__name__}) for sample '
-                                f'{sample_id}.')
-
     def _ensure_autoregressive_input_shapes(self) -> Dict[str, List[int]]:
         handler = global_leap_binder.setup_container.autoregressive_step
         assert handler is not None
         if handler.input_shapes is None:
             preprocess_result = self._preprocess_result()
             first_state_response = next(iter(preprocess_result.values()))
-            first_sample_id = first_state_response.sample_ids[0]
-            first_result = handler.function(first_sample_id, None, None, None,
+            first_result = handler.function(first_state_response.sample_ids[0], None, None,
                                             first_state_response)
-            if not isinstance(first_result, tuple) or len(first_result) != 2 or \
-                    not isinstance(first_result[0], dict):
-                raise Exception('The autoregressive step hook must return a (next_inputs, state) '
-                                'tuple with a dict of initial model inputs on its first call, '
-                                f'got {type(first_result)}.')
-            self._validate_autoregressive_step_inputs(first_result[0], first_sample_id)
-            handler.input_shapes = {key: list(value.shape)
-                                    for key, value in first_result[0].items()}
+            if not isinstance(first_result, dict):
+                raise Exception('The autoregressive step hook must return a dict of initial model '
+                                f'inputs on its first call, got {type(first_result)}.')
+            handler.input_shapes = {key: list(value.shape) for key, value in first_result.items()
+                                    if not key.startswith('_') and isinstance(value, np.ndarray)}
         return handler.input_shapes
 
     def run_simulation(self, sim_name, params=None, n_samples=1, seed=0,
@@ -822,21 +809,6 @@ class LeapLoader(LeapLoaderBase):
             metric_inst = MetricInstance(metric.metric_handler_data.name, metric.metric_handler_data.arg_names)
             metrics.append(metric_inst)
 
-        # Autoregressive decorators are reported alongside the regular ones so the platform sees
-        # their names; the engine tells them apart by name via get_autoregressive_decorator_names.
-        for autoregressive_metric in setup.autoregressive_metrics:
-            metrics.append(MetricInstance(autoregressive_metric.metric_handler_data.name,
-                                          autoregressive_metric.metric_handler_data.arg_names))
-        for autoregressive_loss in setup.autoregressive_losses:
-            custom_losses.append(CustomLossInstance(
-                autoregressive_loss.custom_loss_handler_data.name,
-                autoregressive_loss.custom_loss_handler_data.arg_names))
-        for autoregressive_visualizer in setup.autoregressive_visualizers:
-            visualizers.append(VisualizerInstance(
-                autoregressive_visualizer.visualizer_handler_data.name,
-                autoregressive_visualizer.visualizer_handler_data.type,
-                autoregressive_visualizer.visualizer_handler_data.arg_names))
-
         simulations = []
         for sim in setup.simulations:
             sim_config_serialized = {
@@ -882,6 +854,58 @@ class LeapLoader(LeapLoaderBase):
 
         return sample_ids
 
+    @staticmethod
+    def _grouped_dict_keys(rows: List[Dict[str, Any]], handler_name: str) -> List[str]:
+        """Keys for a group of dict-metadata rows, requiring every row to share the same keys.
+        A group must expose a uniform metadata schema; otherwise the per-key lists would silently
+        drop or KeyError on rows that differ from row 0."""
+        keys = list(rows[0].keys())
+        key_set = set(keys)
+        for pos, row in enumerate(rows):
+            if set(row.keys()) != key_set:
+                raise ValueError(
+                    f"Metadata handler {handler_name!r} returned inconsistent dict keys across the "
+                    f"group: row 0 has {sorted(key_set)} but row {pos} has {sorted(row.keys())}. "
+                    f"All rows in a group must share the same metadata keys.")
+        return keys
+
+    @staticmethod
+    def _to_grouped_list(raw: Any) -> List[npt.NDArray[np.float32]]:
+        """Normalize a grouped encoder result to a list of per-sample arrays.
+
+        Grouped results are NEVER stacked into a (B, *) batch. Grouping is only a storage
+        read-unit (one file load per group); assembling the model batch (B) is the engine's
+        single responsibility, so it stays the only place samples get batched. Keeping
+        per-sample arrays lets the engine batch to any B — including a B=1-only model — without
+        an intermediate stack/unstack. A user encoder that already returned a stacked (B, *)
+        array is split back into its per-sample rows so the contract is a list either way."""
+        if isinstance(raw, list):
+            return raw
+        return [row for row in np.asarray(raw)]
+
+    def _locate_group(self, preprocess_state: PreprocessResponse,
+                      sample_id: Union[int, str]) -> Tuple[List[Union[int, str]], int]:
+        """Map a sample id to (its group, its position in that group). Cached per
+        PreprocessResponse so grouped single-sample fetches don't rescan the groups."""
+        cache = getattr(self, '_group_pos_cache', None)
+        if cache is None:
+            cache = {}
+            self._group_pos_cache = cache
+        mapping = cache.get(id(preprocess_state))
+        if mapping is None:
+            mapping = {}
+            for group in preprocess_state.groups:
+                for pos, sid in enumerate(group):
+                    if sid in mapping:
+                        raise ValueError(
+                            f"Duplicate sample id {sid!r} across groups in the preprocess response; "
+                            f"sample ids must be unique within and across groups.")
+                    mapping[sid] = (group, pos)
+            cache[id(preprocess_state)] = mapping
+        if sample_id not in mapping:
+            raise KeyError(f"Sample id {sample_id!r} is not present in any group of this preprocess response.")
+        return mapping[sample_id]
+
     def _get_dataset_handlers(self, handlers: Iterable[DatasetBaseHandler],
                               state: DataStateEnum, sample_id: Union[int, str]) -> Dict[str, npt.NDArray[np.float32]]:
         result_agg = {}
@@ -893,12 +917,71 @@ class LeapLoader(LeapLoaderBase):
             sample_id = original_local_id
         else:
             preprocess_state = preprocess_result[state]
+        if preprocess_state.is_grouped:
+            # Grouped dataset: call the encoder once with the whole group, index out the sample.
+            group_ids, pos = self._locate_group(preprocess_state, sample_id)
+            for handler in handlers:
+                grouped = self._to_grouped_list(handler.function(group_ids, preprocess_state))
+                result_agg[handler.name] = grouped[pos]
+            return result_agg
         for handler in handlers:
             handler_result = handler.function(sample_id, preprocess_state)
             handler_name = handler.name
             result_agg[handler_name] = handler_result
         return result_agg
 
+    def get_samples(self, state: DataStateEnum, group_ids: List[Union[int, str]]) -> DatasetSample:
+        """Group-aware fetch: hand the whole group to each encoder in a single call and return a
+        grouped DatasetSample (inputs/gt as per-sample lists of length B, metadata as per-key lists
+        of length B, index = the group). Grouped results are never stacked here — batching the
+        samples into the model's B is the engine's job. This is the path a group-aware engine calls
+        to read a file once per group. Row order matches group_ids exactly."""
+        self.exec_script()
+        preprocess_state = self._preprocess_result()[state]
+        if preprocess_state.is_grouped:
+            # All requested ids must belong to a single group (one file). A cross-group request
+            # would force the encoder to load multiple files, defeating the one-load-per-group
+            # contract; partitioning a scattered request by group is the engine's responsibility.
+            distinct_groups = {id(self._locate_group(preprocess_state, sid)[0]) for sid in group_ids}
+            assert len(distinct_groups) == 1, (
+                f"get_samples received sample ids spanning {len(distinct_groups)} groups; a request "
+                f"must be confined to a single group. The engine partitions cross-group requests and "
+                f"issues one get_samples call per group.")
+        inputs = {handler.name: self._to_grouped_list(handler.function(group_ids, preprocess_state))
+                  for handler in global_leap_binder.setup_container.inputs}
+        gt = None
+        if state != DataStateEnum.unlabeled:
+            gt = {handler.name: self._to_grouped_list(handler.function(group_ids, preprocess_state))
+                  for handler in global_leap_binder.setup_container.ground_truths}
+
+        metadata: Dict[str, Any] = {}
+        metadata_is_none: Dict[str, Any] = {}
+        for handler in global_leap_binder.setup_container.metadata:
+            rows = handler.function(group_ids, preprocess_state)  # list of B scalars/dicts
+            if rows and isinstance(rows[0], dict):
+                for key in self._grouped_dict_keys(rows, handler.name):
+                    name = "{}_{}".format(handler.name, key)
+                    converted = [self._convert_metadata_to_correct_type(name, row[key]) for row in rows]
+                    metadata[name] = [v for v, _ in converted]
+                    metadata_is_none[name] = [is_none for _, is_none in converted]
+            else:
+                converted = [self._convert_metadata_to_correct_type(handler.name, row) for row in rows]
+                metadata[handler.name] = [v for v, _ in converted]
+                metadata_is_none[handler.name] = [is_none for _, is_none in converted]
+
+        # custom_latent_space is group-aware like the input/GT encoders: hand it the whole group in a
+        # single call so a file-backed latent fn loads the file once, and normalize to (B, d). This
+        # keeps the mandatory group path non-lossy. instance_masks stay None here: they are a
+        # per-(sample, instance) concern that needs an instance_id the group fetch does not carry.
+        custom_latent_space = None
+        if global_leap_binder.setup_container.custom_latent_space is not None:
+            latent_fn = global_leap_binder.setup_container.custom_latent_space.function
+            custom_latent_space = self._to_grouped_list(latent_fn(group_ids, preprocess_state))
+
+        return DatasetSample(inputs=inputs, gt=gt, metadata=metadata, metadata_is_none=metadata_is_none,
+                             index=list(group_ids), state=state, custom_latent_space=custom_latent_space,
+                             instance_masks=None)
+
     def _get_inputs(self, state: DataStateEnum, sample_id: Union[int, str]) -> Dict[str, npt.NDArray[np.float32]]:
         inputs = self._get_dataset_handlers(global_leap_binder.setup_container.inputs, state, sample_id)
         autoregressive_handler = global_leap_binder.setup_container.autoregressive_step
@@ -907,16 +990,13 @@ class LeapLoader(LeapLoaderBase):
             # the sample's (step-0) model inputs, so every sample-fetching flow (graph validation,
             # analysis, import checks) sees a fully-populated inputs dict.
             preprocess_response = self._preprocess_result()[state]
-            step_zero_result = autoregressive_handler.function(sample_id, None, None, None,
-                                                               preprocess_response)
-            if not isinstance(step_zero_result, tuple) or len(step_zero_result) != 2 or \
-                    not isinstance(step_zero_result[0], dict):
+            step_zero_inputs = autoregressive_handler.function(sample_id, None, None, preprocess_response)
+            if not isinstance(step_zero_inputs, dict):
                 raise Exception(
-                    'The autoregressive step hook must return a (next_inputs, state) tuple with a '
-                    f'dict of initial model inputs on its first call, got {type(step_zero_result)} '
-                    f'for sample {sample_id}.')
-            self._validate_autoregressive_step_inputs(step_zero_result[0], sample_id)
-            inputs.update(step_zero_result[0])
+                    'The autoregressive step hook must return a dict of initial model inputs on its '
+                    f'first call, got {type(step_zero_inputs)} for sample {sample_id}.')
+            inputs.update({key: value for key, value in step_zero_inputs.items()
+                           if not key.startswith('_')})
         return inputs
 
     def _get_instances_masks(self, state: DataStateEnum, sample_id: Union[int, str], instance_id: int) -> Optional[Dict[str, ElementInstance]]:
@@ -1002,12 +1082,18 @@ class LeapLoader(LeapLoaderBase):
             sample_id = original_local_id
         else:
             preprocess_state = preprocess_result[state]
+        group_ids, pos = (None, None)
+        if preprocess_state.is_grouped:
+            group_ids, pos = self._locate_group(preprocess_state, sample_id)
         for handler in global_leap_binder.setup_container.metadata:
             if requested_metadata_names:
                 if not is_metadata_name_starts_with_handler_name(handler):
                     continue
 
-            handler_result = handler.function(sample_id, preprocess_state)
+            if preprocess_state.is_grouped:
+                handler_result = handler.function(group_ids, preprocess_state)[pos]
+            else:
+                handler_result = handler.function(sample_id, preprocess_state)
             if isinstance(handler_result, dict):
                 for single_metadata_name, single_metadata_result in handler_result.items():
                     handler_name = f'{handler.name}_{single_metadata_name}'
@@ -1026,6 +1112,66 @@ class LeapLoader(LeapLoaderBase):
 
         return result_agg, is_none
 
+    def _grouped_metadata_for_group(self, preprocess_state: PreprocessResponse,
+                                    group_ids: List[Union[int, str]],
+                                    requested_metadata_names: Optional[List[str]]
+                                    ) -> Tuple[Dict[str, List[Any]], Dict[str, List[bool]]]:
+        def is_wanted(_handler):
+            if not requested_metadata_names:
+                return True
+            for metadata_name in requested_metadata_names:
+                if metadata_name.startswith(_handler.name + '_') or metadata_name == _handler.name:
+                    return True
+            return False
+
+        results: Dict[str, List[Any]] = {}
+        is_none: Dict[str, List[bool]] = {}
+        for handler in global_leap_binder.setup_container.metadata:
+            if requested_metadata_names and not is_wanted(handler):
+                continue
+            rows = handler.function(group_ids, preprocess_state)
+            if rows and isinstance(rows[0], dict):
+                for key in self._grouped_dict_keys(rows, handler.name):
+                    name = f'{handler.name}_{key}'
+                    if requested_metadata_names and name not in requested_metadata_names:
+                        continue
+                    converted = [self._convert_metadata_to_correct_type(name, row[key]) for row in rows]
+                    results[name] = [v for v, _ in converted]
+                    is_none[name] = [n for _, n in converted]
+            else:
+                name = handler.name
+                if requested_metadata_names and name not in requested_metadata_names:
+                    continue
+                converted = [self._convert_metadata_to_correct_type(name, row) for row in rows]
+                results[name] = [v for v, _ in converted]
+                is_none[name] = [n for _, n in converted]
+        return results, is_none
+
+    def get_metadata_multiple_samples(self, state: DataStateEnum, sample_ids: Union[List[int], List[str]],
+                                      requested_metadata_names: Optional[List[str]] = None
+                                      ) -> Tuple[Dict[str, Union[List[str], List[int], List[bool], List[float]]],
+                                                 Dict[str, List[bool]]]:
+        preprocess_state = self._preprocess_result().get(state)
+        if preprocess_state is None or not preprocess_state.is_grouped:
+            return super().get_metadata_multiple_samples(state, sample_ids, requested_metadata_names)
+
+        sample_id_type = self.get_sample_id_type()
+        aggregated_results: Dict[str, List[Any]] = {}
+        aggregated_is_none: Dict[str, List[bool]] = {}
+        group_cache: Dict[int, Tuple[Dict[str, List[Any]], Dict[str, List[bool]]]] = {}
+        for sample_id in sample_ids:
+            sample_id = sample_id_type(sample_id)
+            group_ids, pos = self._locate_group(preprocess_state, sample_id)
+            key = id(group_ids)
+            if key not in group_cache:
+                group_cache[key] = self._grouped_metadata_for_group(
+                    preprocess_state, group_ids, requested_metadata_names)
+            group_results, group_is_none = group_cache[key]
+            for name, values in group_results.items():
+                aggregated_results.setdefault(name, []).append(values[pos])
+                aggregated_is_none.setdefault(name, []).append(group_is_none[name][pos])
+        return aggregated_results, aggregated_is_none
+
     @lru_cache()
     def get_sample_id_type(self) -> Type:
         preprocess_results = list(self._preprocess_result().values())
@@ -1057,15 +1203,12 @@ class LeapLoader(LeapLoaderBase):
     def run_autoregressive_step(self, sample_id: Union[int, str],
                                 prev_inputs: Optional[Dict[str, npt.NDArray[np.float32]]],
                                 prev_outputs: Optional[Dict[str, npt.NDArray[np.float32]]],
-                                chain_state: Any,
-                                state: DataStateEnum
-                                ) -> Tuple[Optional[Dict[str, npt.NDArray[np.float32]]], Any]:
+                                state: DataStateEnum) -> Optional[Dict[str, npt.NDArray[np.float32]]]:
         """Execute the user's autoregressive step hook against the in-process preprocess result.
 
         This is the engine rollout's per-step entry point (invoked on the generic pod). The first
-        call per chain passes prev_inputs=None, prev_outputs=None, chain_state=None and returns
-        the initial model inputs plus the initial state; a None next_inputs ends the chain — the
-        terminating call still returns the final state.
+        call per chain passes prev_inputs=None, prev_outputs=None and returns the initial model
+        inputs; a None return ends the chain.
         """
         self.exec_script()
         handler = global_leap_binder.setup_container.autoregressive_step
@@ -1073,119 +1216,28 @@ class LeapLoader(LeapLoaderBase):
             raise Exception('run_autoregressive_step was called but the integration has no '
                             'tensorleap_autoregressive_step hook.')
         preprocess_response = self._preprocess_result()[state]
-        result = handler.function(sample_id, prev_inputs, prev_outputs, chain_state,
-                                  preprocess_response)
-        if not isinstance(result, tuple) or len(result) != 2:
-            raise Exception('The autoregressive step hook must return a (next_inputs, state) '
-                            f'tuple, got {type(result)} for sample {sample_id}.')
-        next_inputs, new_chain_state = result
-        validate_autoregressive_state_types(new_chain_state)
-        if next_inputs is None:
-            if prev_inputs is None:
-                raise Exception('The autoregressive step hook returned None for next_inputs on '
-                                f'the first call of the chain (sample {sample_id}) — the first '
-                                'call must return the initial model inputs; None would produce '
-                                'an empty chain with no forward pass.')
-            return None, new_chain_state
-        if not isinstance(next_inputs, dict):
+        result = handler.function(sample_id, prev_inputs, prev_outputs, preprocess_response)
+        if result is None:
+            return None
+        if not isinstance(result, dict):
             raise Exception('The autoregressive step hook must return a dict of model inputs or '
-                            f'None to end the chain, got {type(next_inputs)} for sample '
-                            f'{sample_id}.')
+                            f'None to end the chain, got {type(result)} for sample {sample_id}.')
         # Parse-time validation only exercises step 0; enforce the fixed key-set/shape contract
         # on every step here so drift fails with an actionable message instead of a cryptic
         # engine error on a compiled graph.
         expected_shapes = self._ensure_autoregressive_input_shapes()
-        if set(next_inputs) != set(expected_shapes):
+        model_inputs = {key: value for key, value in result.items() if not key.startswith('_')}
+        if set(model_inputs) != set(expected_shapes):
             raise Exception('The autoregressive step hook must return the same model input keys '
                             f'on every call. Expected {sorted(expected_shapes)}, got '
-                            f'{sorted(next_inputs)} for sample {sample_id}.')
-        for key, value in next_inputs.items():
-            if not isinstance(value, np.ndarray):
-                raise Exception('The autoregressive step hook must return numpy arrays for every '
-                                f'model input on every call — input "{key}" is a '
-                                f'{type(value).__name__} for sample {sample_id} (forgot .numpy() '
-                                'on a fed-back model output?).')
-            if list(value.shape) != expected_shapes[key]:
+                            f'{sorted(model_inputs)} for sample {sample_id}.')
+        for key, value in model_inputs.items():
+            if isinstance(value, np.ndarray) and list(value.shape) != expected_shapes[key]:
                 raise Exception('The autoregressive step hook must return fixed tensor shapes on '
                                 f'every call (use fixed-length padding + mask for growing '
                                 f'sequences). Input "{key}" was {expected_shapes[key]} on the '
                                 f'first call and {list(value.shape)} now (sample {sample_id}).')
-        return next_inputs, new_chain_state
-
-    @lru_cache()
-    def autoregressive_visualizer_by_name(self) -> Dict[str, VisualizerHandlerData]:
-        self.exec_script()
-        return {handler.visualizer_handler_data.name: handler.visualizer_handler_data
-                for handler in global_leap_binder.setup_container.autoregressive_visualizers}
-
-    @lru_cache()
-    def get_autoregressive_decorator_names(self) -> Dict[str, List[str]]:
-        self.exec_script()
-        setup = global_leap_binder.setup_container
-        return {
-            'metrics': [handler.metric_handler_data.name
-                        for handler in setup.autoregressive_metrics],
-            'losses': [handler.custom_loss_handler_data.name
-                       for handler in setup.autoregressive_losses],
-            'visualizers': [handler.visualizer_handler_data.name
-                            for handler in setup.autoregressive_visualizers],
-        }
-
-    @staticmethod
-    def _autoregressive_callable_kwargs(function: Callable[..., Any],
-                                        inputs: Dict[str, npt.NDArray[np.float32]],
-                                        outputs: Dict[str, npt.NDArray[np.float32]],
-                                        chain_state: Any,
-                                        wired_args: Dict[str, npt.NDArray[np.float32]]
-                                        ) -> Dict[str, Any]:
-        implicit_values = {'inputs': inputs, 'outputs': outputs, 'state': chain_state}
-        kwargs: Dict[str, Any] = {}
-        for arg_name in inspect.getfullargspec(function)[0]:
-            if arg_name in implicit_values:
-                kwargs[arg_name] = implicit_values[arg_name]
-            elif arg_name in wired_args:
-                kwargs[arg_name] = wired_args[arg_name]
-            else:
-                raise Exception(f'Autoregressive callable is missing wired argument '
-                                f'"{arg_name}" — the platform mapping supplies only '
-                                f'{sorted(wired_args)}.')
-        return kwargs
-
-    def run_autoregressive_metric(self, name: str,
-                                  inputs: Dict[str, npt.NDArray[np.float32]],
-                                  outputs: Dict[str, npt.NDArray[np.float32]],
-                                  chain_state: Any,
-                                  wired_args: Dict[str, npt.NDArray[np.float32]]) -> Any:
-        self.exec_script()
-        for handler in global_leap_binder.setup_container.autoregressive_metrics:
-            if handler.metric_handler_data.name == name:
-                return handler.function(**self._autoregressive_callable_kwargs(
-                    handler.function, inputs, outputs, chain_state, wired_args))
-        raise Exception(f'No autoregressive metric named "{name}" is registered.')
-
-    def run_autoregressive_loss(self, name: str,
-                                inputs: Dict[str, npt.NDArray[np.float32]],
-                                outputs: Dict[str, npt.NDArray[np.float32]],
-                                chain_state: Any,
-                                wired_args: Dict[str, npt.NDArray[np.float32]]) -> Any:
-        self.exec_script()
-        for handler in global_leap_binder.setup_container.autoregressive_losses:
-            if handler.custom_loss_handler_data.name == name:
-                return handler.function(**self._autoregressive_callable_kwargs(
-                    handler.function, inputs, outputs, chain_state, wired_args))
-        raise Exception(f'No autoregressive loss named "{name}" is registered.')
-
-    def run_autoregressive_visualizer(self, name: str,
-                                      inputs: Dict[str, npt.NDArray[np.float32]],
-                                      outputs: Dict[str, npt.NDArray[np.float32]],
-                                      chain_state: Any,
-                                      wired_args: Dict[str, npt.NDArray[np.float32]]) -> Any:
-        self.exec_script()
-        for handler in global_leap_binder.setup_container.autoregressive_visualizers:
-            if handler.visualizer_handler_data.name == name:
-                return handler.function(**self._autoregressive_callable_kwargs(
-                    handler.function, inputs, outputs, chain_state, wired_args))
-        raise Exception(f'No autoregressive visualizer named "{name}" is registered.')
+        return result
 
     def get_prediction_names_in_order(self) -> List[str]:
         """Model output names by declared prediction-type order, used to key prev_outputs for the