code-loader 1.0.195.dev1__py3-none-any.whl → 1.0.197.dev0__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, Union, Any, Type, Optional, Callable, Tuple
+from typing import Dict, List, Iterable, Set, FrozenSet, Union, Any, Type, Optional, Callable, Tuple
 
 import numpy as np
 import numpy.typing as npt
@@ -27,7 +27,8 @@ 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
+from code_loader.utils import get_root_exception_file_and_line_number, get_metadata_type_from_variable, \
+    validate_autoregressive_state_types, autoregressive_nests_equal
 
 
 def _serialize_sim_bounds(bounds) -> dict:
@@ -59,35 +60,14 @@ 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:
-                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)
+                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))
         except TypeError as e:
             import traceback
             global_leap_binder.setup_container = DatasetIntegrationSetup()
@@ -100,7 +80,6 @@ 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]
 
@@ -271,20 +250,12 @@ class LeapLoader(LeapLoaderBase):
             )
 
         preprocess_result = self._preprocess_result()
-        if state == DataStateEnum.unlabeled and sample_id not in preprocess_result[state].flat_sample_ids:
+        if state == DataStateEnum.unlabeled and sample_id not in preprocess_result[state].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:
-            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)
+        custom_latent_spaces = self._get_custom_latent_spaces(sample_id, preprocess_result[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),
@@ -292,7 +263,7 @@ class LeapLoader(LeapLoaderBase):
                                metadata_is_none=metadata_is_none,
                                index=sample_id,
                                state=state,
-                               custom_latent_space=custom_latent_space,
+                               custom_latent_spaces=custom_latent_spaces,
                                instance_masks=instance_mask)
         return sample
 
@@ -323,11 +294,7 @@ class LeapLoader(LeapLoaderBase):
                     handler.name, handler_result
                 )
 
-        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(
-                original_sample_id, preprocess
-            )
+        custom_latent_spaces = self._get_custom_latent_spaces(original_sample_id, preprocess)
 
         return DatasetSample(
             inputs=inputs,
@@ -336,7 +303,7 @@ class LeapLoader(LeapLoaderBase):
             metadata_is_none=metadata_is_none,
             index=synthetic_index,
             state=DataStateEnum.additional,
-            custom_latent_space=custom_latent_space,
+            custom_latent_spaces=custom_latent_spaces,
             instance_masks=None,
         )
 
@@ -415,7 +382,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.flat_sample_ids)
+                    sample_ids_array = np.array(preprocess_response.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.")
 
@@ -505,31 +472,49 @@ 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, preprocess_response)
-                second_result = handler.function(sample_id, None, None, preprocess_response)
-                if not isinstance(first_result, dict):
+                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:
                     raise Exception('The autoregressive step hook must return a dict of initial model '
-                                    f'inputs on its first call, got {type(first_result)} '
+                                    f'inputs on its first call, got {type(first_inputs)} '
                                     f'(state: {state.name}).')
-                if not isinstance(second_result, dict) or \
-                        set(second_result.keys()) != set(first_result.keys()):
+                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]):
                     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.')
-                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 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}.')
                     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}.')
@@ -538,7 +523,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 np.array_equal(value, second_result[key]):
+                    if not autoregressive_nests_equal(value, second_result[0][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 '
@@ -552,19 +537,39 @@ 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_result = handler.function(first_state_response.sample_ids[0], None, None,
+            first_sample_id = first_state_response.sample_ids[0]
+            first_result = handler.function(first_sample_id, None, None, None,
                                             first_state_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_result)}.')
-            handler.input_shapes = {key: list(value.shape) for key, value in first_result.items()
-                                    if not key.startswith('_') and isinstance(value, np.ndarray)}
+            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()}
         return handler.input_shapes
 
     def run_simulation(self, sim_name, params=None, n_samples=1, seed=0,
@@ -833,6 +838,21 @@ 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 = {
@@ -878,58 +898,6 @@ 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 = {}
@@ -941,71 +909,12 @@ 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
@@ -1014,13 +923,16 @@ 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_inputs = autoregressive_handler.function(sample_id, None, None, preprocess_response)
-            if not isinstance(step_zero_inputs, dict):
+            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):
                 raise Exception(
-                    '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('_')})
+                    '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])
         return inputs
 
     def _get_instances_masks(self, state: DataStateEnum, sample_id: Union[int, str], instance_id: int) -> Optional[Dict[str, ElementInstance]]:
@@ -1106,18 +1018,12 @@ 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
 
-            if preprocess_state.is_grouped:
-                handler_result = handler.function(group_ids, preprocess_state)[pos]
-            else:
-                handler_result = handler.function(sample_id, preprocess_state)
+            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}'
@@ -1136,66 +1042,6 @@ 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())
@@ -1206,10 +1052,29 @@ class LeapLoader(LeapLoaderBase):
 
         return id_type
 
+    @staticmethod
+    def _get_custom_latent_spaces(
+            sample_id: Union[int, str],
+            preprocess: "PreprocessResponse") -> Optional[Dict[str, npt.NDArray[np.float32]]]:
+        handlers = global_leap_binder.setup_container.custom_latent_spaces
+        if not handlers:
+            return None
+        return {name: handler.function(sample_id, preprocess) for name, handler in handlers.items()}
+
     @lru_cache()
     def has_custom_latent_space_decorator(self) -> bool:
         self.exec_script()
-        return global_leap_binder.setup_container.custom_latent_space is not None
+        return len(global_leap_binder.setup_container.custom_latent_spaces) > 0
+
+    @lru_cache()
+    def get_custom_latent_space_names(self) -> Tuple[str, ...]:
+        """Ordered names of all registered custom latent spaces.
+
+        Registration order is the canonical index mapping consumed by the engine
+        (name i -> user_custom_i). Returns a tuple so the lru_cache value is hashable.
+        """
+        self.exec_script()
+        return tuple(global_leap_binder.setup_container.custom_latent_spaces.keys())
 
     @lru_cache()
     def has_autoregressive_step(self) -> bool:
@@ -1227,12 +1092,15 @@ 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]]],
-                                state: DataStateEnum) -> Optional[Dict[str, npt.NDArray[np.float32]]]:
+                                chain_state: Any,
+                                state: DataStateEnum
+                                ) -> Tuple[Optional[Dict[str, npt.NDArray[np.float32]]], Any]:
         """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 and returns the initial model
-        inputs; a None return ends the chain.
+        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.
         """
         self.exec_script()
         handler = global_leap_binder.setup_container.autoregressive_step
@@ -1240,28 +1108,119 @@ 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, preprocess_response)
-        if result is None:
-            return None
-        if not isinstance(result, dict):
+        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):
             raise Exception('The autoregressive step hook must return a dict of model inputs or '
-                            f'None to end the chain, got {type(result)} for sample {sample_id}.')
+                            f'None to end the chain, got {type(next_inputs)} for sample '
+                            f'{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()
-        model_inputs = {key: value for key, value in result.items() if not key.startswith('_')}
-        if set(model_inputs) != set(expected_shapes):
+        if set(next_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(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]:
+                            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]:
                 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 result
+        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.')
 
     def get_prediction_names_in_order(self) -> List[str]:
         """Model output names by declared prediction-type order, used to key prev_outputs for the