code-loader 1.0.191.dev1__py3-none-any.whl → 1.0.193.dev0__py3-none-any.whl

code_loader/inner_leap_binder/leapbinder_decorators.py

@@ -1,4 +1,6 @@
 # mypy: ignore-errors
+import copy
+import hashlib
 import inspect
 import os
 import warnings
@@ -13,7 +15,8 @@ from typing import Optional, Union, Callable, List, Dict, get_args, get_origin,
 import numpy as np
 import numpy.typing as npt
 
-from code_loader.utils import get_metadata_type_from_variable, map_dict_to_metadata_types
+from code_loader.utils import get_metadata_type_from_variable, map_dict_to_metadata_types, \
+    validate_autoregressive_state_types, autoregressive_nests_equal
 
 logger = logging.getLogger(__name__)
 
@@ -21,18 +24,18 @@ from code_loader.contract.datasetclasses import CustomCallableInterfaceMultiArgs
     CustomMultipleReturnCallableInterfaceMultiArgs, ConfusionMatrixCallableInterfaceMultiArgs, CustomCallableInterface, \
     VisualizerCallableInterface, MetadataSectionCallableInterface, PreprocessResponse, SectionCallableInterface, \
     ConfusionMatrixElement, SamplePreprocessResponse, PredictionTypeHandler, InstanceCallableInterface, ElementInstance, \
-    InstanceLengthCallableInterface
+    InstanceLengthCallableInterface, AutoregressiveStepCallableInterface
 from code_loader.contract.enums import MetricDirection, LeapDataType, DatasetMetadataType, DataStateType
 from code_loader import leap_binder, LeapLoader
 from code_loader.contract.mapping import NodeMapping, NodeMappingType, NodeConnection
 from code_loader.contract.visualizer_classes import LeapImage, LeapImageMask, LeapTextMask, LeapText, LeapGraph, \
-    LeapHorizontalBar, LeapImageWithBBox, LeapImageWithHeatmap, LeapVideo
+    LeapHorizontalBar, LeapImageWithBBox, LeapImageWithHeatmap, LeapVideo, LeapValidationError, \
+    map_leap_data_type_to_visualizer_class
 from code_loader.inner_leap_binder.leapbinder import mapping_runtime_mode_env_var_mame
 from code_loader.mixpanel_tracker import clear_integration_events, AnalyticsEvent, emit_integration_event_once
 
 _called_from_inside_tl_decorator = 0
 _called_from_inside_tl_integration_test_decorator = False
-_mapping_dataset_is_grouped = False
 _call_from_tl_platform = os.environ.get('IS_TENSORLEAP_PLATFORM') == 'true'
 
 # ---- warnings store (module-level) ----
@@ -198,58 +201,14 @@ def validate_output_structure(result, func_name: str, expected_type_name="np.nda
 def batch_warning(result, func_name):
     if len(result.shape) > 0 and result.shape[0] == 1:
         warnings.warn(
-            f"{func_name} warning: Tensorleap will add a batch dimension at axis 0 to the output of {func_name}, "
-            f"although the detected size of axis 0 is already 1. "
-            f"This may lead to an extra batch dimension (e.g., shape (1, 1, ...)). "
-            f"Please ensure that the output of '{func_name}' is not already batched "
-            f"to avoid computation errors."
+            f"{func_name} warning: Tensorleap will add a batch dimension at axis 0 of "
+            f"{func_name}'s output, but axis 0 is already size 1 — this can produce a "
+            f"doubly-batched shape (e.g. (1, 1, ...)). If axis 0 is a real dimension "
+            f"(e.g. a single channel) and not an accidental batch, you can ignore this; "
+            f"only fix it if you pre-batched the output."
         )
 
 
-def _validate_id_or_group(sample_id, preprocess_response, func_name):
-    """Accept either a single sample id (flat dataset) or a group — a list of ids (grouped
-    dataset). The argument shape must match the dataset: a group requires a grouped
-    PreprocessResponse and a single id requires a flat one. Every id must match the declared
-    sample_id_type."""
-    is_group = isinstance(sample_id, list)
-    if is_group and not preprocess_response.is_grouped:
-        raise AssertionError(
-            f'{func_name}() validation failed: got a group (list of sample ids) but the '
-            f'PreprocessResponse is not grouped. Pass a single sample id for a flat dataset.')
-    if not is_group and preprocess_response.is_grouped:
-        raise AssertionError(
-            f'{func_name}() validation failed: got a single sample id but the PreprocessResponse '
-            f'is grouped. Pass a group (list of sample ids) for a grouped dataset.')
-    ids = sample_id if is_group else [sample_id]
-    for sid in ids:
-        assert type(sid) == preprocess_response.sample_id_type, \
-            (f'{func_name}() validation failed: '
-             f'Argument sample_id should be as the same type as defined in the preprocess response '
-             f'{preprocess_response.sample_id_type}. Got {type(sid)}.')
-
-
-def _validate_grouped_result(result, group_size, func_name, validate_single):
-    """A grouped (group) encoder result must be either a single ndarray with leading dim
-    == group_size, or a list of `group_size` per-sample results (each validated by
-    validate_single). Metadata validates its grouped result separately (list only)."""
-    if isinstance(result, np.ndarray):
-        assert len(result.shape) > 0 and result.shape[0] == group_size, \
-            (f'{func_name}() validation failed: expected a grouped output for a group of '
-             f'{group_size}: an array with leading dim {group_size}, got shape {result.shape}.')
-        validate_single(result)
-    elif isinstance(result, list):
-        assert len(result) == group_size, \
-            (f'{func_name}() validation failed: expected a grouped output for a group of '
-             f'{group_size}: a list of {group_size} arrays, got {len(result)}.')
-        for r in result:
-            validate_single(r)
-    else:
-        raise AssertionError(
-            f'{func_name}() validation failed: expected a grouped output for a group of '
-            f'{group_size}: either an array with leading dim {group_size} or a list of '
-            f'{group_size} arrays, got {type(result)}.')
-
-
 def _add_mapping_connection(user_unique_name, connection_destinations, arg_names, name, node_mapping_type):
     connection_destinations = [connection_destination for connection_destination in connection_destinations
                                if not isinstance(connection_destination, SamplePreprocessResponse)]
@@ -327,12 +286,16 @@ def tensorleap_integration_test():
 
         def _validate_input_args(*args, **kwargs):
             sample_id, preprocess_response = args
-            _validate_id_or_group(sample_id, preprocess_response, 'tensorleap_integration_test')
+            assert type(sample_id) == preprocess_response.sample_id_type, (
+                f"tensorleap_integration_test validation failed: "
+                f"sample_id type ({type(sample_id).__name__}) does not match the expected "
+                f"type ({preprocess_response.sample_id_type}) from the PreprocessResponse."
+            )
 
         def inner(*args, **kwargs):
             if not _call_from_tl_platform:
                 set_current('tensorleap_integration_test')
-            validate_args_structure(*args, types_order=[Union[int, str, list], PreprocessResponse],
+            validate_args_structure(*args, types_order=[Union[int, str], PreprocessResponse],
                                     func_name='integration_test', expected_names=["idx", "preprocess"], **kwargs)
             _validate_input_args(*args, **kwargs)
 
@@ -347,30 +310,29 @@ def tensorleap_integration_test():
                 if not _call_from_tl_platform:
                     update_env_params_func("tensorleap_integration_test",
                                            "v")  # put here because otherwise it will become  v only if it finishes all the script
+                _reset_model_loop_state()
                 ret = integration_test_function(*args, **kwargs)
 
-                global _mapping_dataset_is_grouped
-                _mapping_dataset_is_grouped = args[1].is_grouped
                 try:
                     os.environ[mapping_runtime_mode_env_var_mame] = 'True'
+                    _reset_model_loop_state()
                     integration_test_function(None, PreprocessResponse(state=DataStateType.training, length=0))
                 except Exception as e:
                     import traceback
                     first_tb = traceback.extract_tb(e.__traceback__)[-1]
                     file_name = Path(first_tb.filename).name
                     line_number = first_tb.lineno
-                    if isinstance(e, TypeError) and 'is not subscriptable' in str(e):
-                        update_env_params_func("code_mapping", "x")
+                    update_env_params_func("code_mapping", "x")
+                    if isinstance(e, LeapValidationError):
+                        raise Exception(f'Invalid integration code. File {file_name}, line {line_number}: {e}')
+                    elif isinstance(e, TypeError) and 'is not subscriptable' in str(e):
                         raise Exception(f'Invalid integration code. File {file_name}, line {line_number}: '
-                               f"indexing is supported only on the model's predictions inside the integration test. Please remove this indexing operation usage from the integration test code.")
+                                        f"indexing is supported only on the model's predictions inside the integration test. Please remove this indexing operation usage from the integration test code.")
                     else:
-                        update_env_params_func("code_mapping", "x")
-
                         raise Exception(f'Invalid integration code. File {file_name}, line {line_number}: '
-                               f'Integration test is only allowed to call Tensorleap decorators. '
-                               f'Ensure any arithmetics, external library use, Python logic is placed within Tensorleap decoders')
+                                        f'Integration test is only allowed to call Tensorleap decorators. '
+                                        f'Ensure any arithmetics, external library use, Python logic is placed within Tensorleap decoders')
                 finally:
-                    _mapping_dataset_is_grouped = False
                     if mapping_runtime_mode_env_var_mame in os.environ:
                         del os.environ[mapping_runtime_mode_env_var_mame]
             finally:
@@ -393,6 +355,14 @@ def _safe_get_item(key):
 
 
 def tensorleap_load_model(prediction_types: Optional[List[PredictionTypeHandler]] = _UNSET):
+    """Register the model-loading function.
+
+    ``prediction_types``: declare one ``PredictionTypeHandler`` per model output, in
+    output order. Multi-output models must declare an entry for EVERY output — the
+    count is validated against the model's outputs — so for an output you don't consume
+    as a prediction, pass a throwaway handler to satisfy the count. Omit the argument
+    only when the model has a single prediction output.
+    """
     prediction_types_was_provided = prediction_types is not _UNSET
 
     if not prediction_types_was_provided:
@@ -612,6 +582,8 @@ def tensorleap_load_model(prediction_types: Optional[List[PredictionTypeHandler]
                     return self.model.get_inputs()
 
             model_placeholder = ModelPlaceholder(prediction_types)
+            global _last_loaded_model
+            _last_loaded_model = model_placeholder
             if not _call_from_tl_platform:
                 update_env_params_func("tensorleap_load_model", "v")
             return model_placeholder
@@ -622,8 +594,16 @@ def tensorleap_load_model(prediction_types: Optional[List[PredictionTypeHandler]
                     self.node_mapping = NodeMapping('', NodeMappingType.Prediction0)
 
                 def __getitem__(self, key):
-                    assert isinstance(key, int), \
-                        f'Expected key to be an int, got {type(key)} instead.'
+                    assert isinstance(key, (int, str)), \
+                        f'Expected key to be an int or a prediction-type name, got {type(key)} instead.'
+                    if isinstance(key, str):
+                        declared_names = [prediction_type.name for prediction_type
+                                          in leap_binder.setup_container.prediction_types]
+                        if key not in declared_names:
+                            raise Exception(
+                                f'Unknown prediction name "{key}" — declared prediction types: '
+                                f'{declared_names}.')
+                        key = declared_names.index(key)
 
                     ret = TempMapping()
                     try:
@@ -647,7 +627,9 @@ def tensorleap_load_model(prediction_types: Optional[List[PredictionTypeHandler]
 
                 # onnx runtime interface
                 def run(self, output_names, input_dict):
-                    assert output_names is None
+                    if output_names is not None:
+                        raise LeapValidationError("Tensorleap doesn't support selecting outputs by name — "
+                                                  "use model.run(None, inputs) to return all outputs.")
                     assert isinstance(input_dict, dict), \
                         f'Expected input_dict to be a dict, got {type(input_dict)} instead.'
                     for i, (input_key, elem) in enumerate(input_dict.items()):
@@ -675,7 +657,10 @@ def tensorleap_load_model(prediction_types: Optional[List[PredictionTypeHandler]
 
                     return FollowInputIndex()
 
-            return ModelPlaceholder()
+            model_placeholder = ModelPlaceholder()
+            global _last_loaded_model
+            _last_loaded_model = model_placeholder
+            return model_placeholder
 
         def final_inner(*args, **kwargs):
             if os.environ.get(mapping_runtime_mode_env_var_mame):
@@ -1552,10 +1537,13 @@ def tensorleap_metadata(
                 raise Exception(f'Metadata with name {name} already exists. '
                                 f'Please choose another')
 
-        def _validate_input_args(sample_id: Union[int, str, list], preprocess_response: PreprocessResponse):
-            _validate_id_or_group(sample_id, preprocess_response, user_function.__name__)
+        def _validate_input_args(sample_id: Union[int, str], preprocess_response: PreprocessResponse):
+            assert type(sample_id) == preprocess_response.sample_id_type, \
+                (f'{user_function.__name__}() validation failed: '
+                 f'Argument sample_id should be as the same type as defined in the preprocess response '
+                 f'{preprocess_response.sample_id_type}. Got {type(sample_id)}.')
 
-        def _validate_single(result):
+        def _validate_result(result):
             supported_result_types = (type(None), int, str, bool, float, dict, np.floating,
                                       np.bool_, np.unsignedinteger, np.signedinteger, np.integer)
             if isinstance(result, tuple):
@@ -1573,20 +1561,6 @@ def tensorleap_metadata(
                         (f'{user_function.__name__}() validation failed: '
                          f'Values in the return dict should be of type {str(supported_result_types)}. Got {type(value)}.')
 
-        def _validate_result(result, grouped=False, group_size=None):
-            if not grouped:
-                _validate_single(result)
-                return
-            # Grouped metadata: a list of `group_size` per-sample scalars/dicts.
-            assert isinstance(result, list), \
-                (f'{user_function.__name__}() validation failed: expected a grouped output for a group '
-                 f'of {group_size}: a list of {group_size} metadata values, got {type(result)}.')
-            assert len(result) == group_size, \
-                (f'{user_function.__name__}() validation failed: expected a grouped output for a group '
-                 f'of {group_size}: a list of {group_size} metadata values, got {len(result)}.')
-            for value in result:
-                _validate_single(value)
-
         def inner_without_validate(sample_id, preprocess_response):
 
             global _called_from_inside_tl_decorator
@@ -1606,16 +1580,14 @@ def tensorleap_metadata(
                 set_current('tensorleap_metadata')
             if os.environ.get(mapping_runtime_mode_env_var_mame):
                 return None
-            validate_args_structure(*args, types_order=[Union[int, str, list], PreprocessResponse],
+            validate_args_structure(*args, types_order=[Union[int, str], PreprocessResponse],
                                     func_name=user_function.__name__, expected_names=["idx", "preprocess"], **kwargs)
             sample_id, preprocess_response = args if len(args) != 0 else kwargs.values()
             _validate_input_args(sample_id, preprocess_response)
 
-            grouped = isinstance(sample_id, list)
-            group_size = len(sample_id) if grouped else None
             result = inner_without_validate(sample_id, preprocess_response)
 
-            _validate_result(result, grouped, group_size)
+            _validate_result(result)
             if not _call_from_tl_platform:
                 update_env_params_func("tensorleap_metadata", "v")
             return result
@@ -1627,50 +1599,35 @@ def tensorleap_metadata(
 
 def tensorleap_custom_latent_space():
     def decorating_function(user_function: SectionCallableInterface):
-        def _validate_input_args(sample_id: Union[int, str, list], preprocess_response: PreprocessResponse):
-            _validate_id_or_group(sample_id, preprocess_response, 'tensorleap_custom_latent_space')
+        def _validate_input_args(sample_id: Union[int, str], preprocess_response: PreprocessResponse):
+            assert isinstance(sample_id, (int, str)), \
+                (f'tensorleap_custom_latent_space validation failed: '
+                 f'Argument sample_id should be either int or str. Got {type(sample_id)}.')
+            assert isinstance(preprocess_response, PreprocessResponse), \
+                (f'tensorleap_custom_latent_space validation failed: '
+                 f'Argument preprocess_response should be a PreprocessResponse. Got {type(preprocess_response)}.')
+            assert type(sample_id) == preprocess_response.sample_id_type, \
+                (f'tensorleap_custom_latent_space validation failed: '
+                 f'Argument sample_id should be as the same type as defined in the preprocess response '
+                 f'{preprocess_response.sample_id_type}. Got {type(sample_id)}.')
 
-        def _validate_single(single_result):
-            assert isinstance(single_result, np.ndarray), \
+        def _validate_result(result):
+            assert isinstance(result, np.ndarray), \
                 (f'tensorleap_custom_latent_space validation failed: '
-                 f'The return type should be a numpy array. Got {type(single_result)}.')
-            if single_result.ndim > 1:
-                flat_dim = int(np.prod(single_result.shape))
+                 f'The return type should be a numpy array. Got {type(result)}.')
+            if result.ndim > 1:
+                flat_dim = int(np.prod(result.shape))
                 store_general_warning(
-                    key=("tensorleap_custom_latent_space_flatten", tuple(single_result.shape)),
+                    key=("tensorleap_custom_latent_space_flatten", tuple(result.shape)),
                     message=(
-                        f"tensorleap_custom_latent_space returned per-sample shape {tuple(single_result.shape)} "
-                        f"(ndim={single_result.ndim}). Tensorleap assumes per-sample shape (d, ...) and will "
+                        f"tensorleap_custom_latent_space returned per-sample shape {tuple(result.shape)} "
+                        f"(ndim={result.ndim}). Tensorleap assumes per-sample shape (d, ...) and will "
                         f"flatten to ({flat_dim},) before downstream visualization and clustering. "
                         f"If you want a different aggregation (e.g. global average pooling), do it "
                         f"inside your function."
                     ),
                 )
 
-        def _validate_result(result, grouped=False, group_size=None):
-            if not grouped:
-                _validate_single(result)
-                return
-            # Grouped: (B, *) array or list of B per-sample arrays. Validate each per-sample slice so
-            # the per-sample flatten check/warning uses the sample shape, not the (B, *) batch shape.
-            if isinstance(result, np.ndarray):
-                assert result.ndim >= 1 and result.shape[0] == group_size, \
-                    (f'tensorleap_custom_latent_space validation failed: expected a grouped output for a '
-                     f'group of {group_size}: an array with leading dim {group_size}, got shape {result.shape}.')
-                for single in result:
-                    _validate_single(single)
-            elif isinstance(result, list):
-                assert len(result) == group_size, \
-                    (f'tensorleap_custom_latent_space validation failed: expected a grouped output for a '
-                     f'group of {group_size}: a list of {group_size} arrays, got {len(result)}.')
-                for single in result:
-                    _validate_single(single)
-            else:
-                raise AssertionError(
-                    f'tensorleap_custom_latent_space validation failed: expected a grouped output for a '
-                    f'group of {group_size}: either an array with leading dim {group_size} or a list of '
-                    f'{group_size} arrays, got {type(result)}.')
-
         def inner_without_validate(sample_id, preprocess_response):
             global _called_from_inside_tl_decorator
             _called_from_inside_tl_decorator += 1
@@ -1690,11 +1647,9 @@ def tensorleap_custom_latent_space():
 
             _validate_input_args(sample_id, preprocess_response)
 
-            grouped = isinstance(sample_id, list)
-            group_size = len(sample_id) if grouped else None
             result = inner_without_validate(sample_id, preprocess_response)
 
-            _validate_result(result, grouped, group_size)
+            _validate_result(result)
             return result
 
         return inner
@@ -1702,6 +1657,939 @@ def tensorleap_custom_latent_space():
     return decorating_function
 
 
+_MODEL_LOOP_MAX_STEPS = 1000
+
+_active_model_loop = None
+_model_loop_run_in_current_test = False
+_last_loaded_model = None
+_chain_artifact_provenance: Dict[int, tuple] = {}
+
+_validate_state_types = validate_autoregressive_state_types
+_nest_equal = autoregressive_nests_equal
+
+
+def _nest_fingerprint(value):
+    digest = hashlib.sha1()
+
+    def _feed(item):
+        if isinstance(item, dict):
+            digest.update(b'{')
+            for key in sorted(item):
+                digest.update(str(key).encode())
+                digest.update(b':')
+                _feed(item[key])
+            digest.update(b'}')
+        elif isinstance(item, (list, tuple)):
+            digest.update(b'[')
+            for element in item:
+                _feed(element)
+            digest.update(b']')
+        elif isinstance(item, np.ndarray):
+            digest.update(str(item.dtype).encode())
+            digest.update(str(item.shape).encode())
+            digest.update(np.ascontiguousarray(item).tobytes())
+        else:
+            digest.update(repr(item).encode())
+
+    _feed(value)
+    return digest.hexdigest()
+
+
+def _ndarray_leaves(value):
+    if isinstance(value, np.ndarray):
+        return [value]
+    if isinstance(value, dict):
+        return [leaf for item in value.values() for leaf in _ndarray_leaves(item)]
+    if isinstance(value, (list, tuple)):
+        return [leaf for item in value for leaf in _ndarray_leaves(item)]
+    return []
+
+
+def _squeeze_leading_batch_axis(tensors_dict):
+    squeezed = {}
+    for key, value in tensors_dict.items():
+        if isinstance(value, np.ndarray) and value.ndim > 0 and value.shape[0] == 1:
+            squeezed[key] = value[0]
+        else:
+            squeezed[key] = value
+    return squeezed
+
+
+def _register_chain_artifact(obj, kind):
+    # The object reference is kept in the registry entry so its id() cannot be recycled by the GC.
+    _chain_artifact_provenance[id(obj)] = (kind, _nest_fingerprint(obj), obj)
+
+
+def _lookup_chain_artifact(obj):
+    entry = _chain_artifact_provenance.get(id(obj))
+    if entry is None or entry[2] is not obj:
+        return None
+    return entry
+
+
+def _reset_model_loop_state():
+    global _active_model_loop, _model_loop_run_in_current_test
+    _active_model_loop = None
+    _model_loop_run_in_current_test = False
+    _chain_artifact_provenance.clear()
+
+
+class _MappingStatePlaceholder:
+    def _unsupported(self):
+        raise LeapValidationError(
+            'The autoregressive state is opaque inside the integration test — pass it whole '
+            'between the hook, the model loop and the autoregressive metrics. Read or derive '
+            'state content inside the hook or inside the metric/visualizer body.')
+
+    def __getitem__(self, key):
+        self._unsupported()
+
+    def __getattr__(self, name):
+        raise AttributeError(name)
+
+    def __iter__(self):
+        self._unsupported()
+
+
+class _ModelLoopContext:
+    def __init__(self, sample_id, preprocess_response, prediction_names, is_mapping):
+        self.sample_id = sample_id
+        self.preprocess_response = preprocess_response
+        self.prediction_names = prediction_names
+        self.is_mapping = is_mapping
+        self.phase = 'awaiting_first_hook'
+        self.steps = 0
+        self.fed_inputs = None
+        self.fed_inputs_fingerprint = None
+        self.last_outputs = None
+        self.last_outputs_fingerprint = None
+        self.last_state = None
+        self.last_state_fingerprint = None
+
+    def _fail(self, message):
+        raise LeapValidationError(f'tensorleap_model_loop validation failed: {message}')
+
+    def on_hook_call(self, sample_id, prev_inputs, prev_outputs, state, preprocess_response):
+        if self.phase == 'ended':
+            self._fail('the chain already ended — once the autoregressive step returns None for '
+                       'next_inputs the loop must exit without further hook or model calls.')
+        if self.phase == 'awaiting_model':
+            self._fail('two autoregressive step calls in a row — the model must run exactly once '
+                       'between hook calls.')
+        if sample_id != self.sample_id:
+            self._fail('the hook must be called with the sample_id the model loop received.')
+        if preprocess_response is not self.preprocess_response:
+            self._fail('the hook must be called with the preprocess response the model loop '
+                       'received.')
+        if self.phase == 'awaiting_first_hook':
+            if prev_inputs is not None or prev_outputs is not None or state is not None:
+                self._fail('the first hook call inside the model loop must be the initial call: '
+                           'prev_inputs=None, prev_outputs=None, state=None.')
+            return
+        if prev_inputs is not self.fed_inputs:
+            self._fail('prev_inputs must be exactly the inputs dict the previous hook call '
+                       'returned (the one fed to the model) — do not copy, rebuild or slice it.')
+        if prev_outputs is not self.last_outputs:
+            self._fail('prev_outputs must be exactly the outputs dict the model call returned — '
+                       'do not copy, rebuild or slice it.')
+        if state is not self.last_state:
+            self._fail('state must be exactly the state object the previous hook call returned.')
+        if _nest_fingerprint(prev_inputs) != self.fed_inputs_fingerprint:
+            self._fail('the model inputs were mutated in place after being fed to the model — '
+                       'the platform re-reads the stored step tensors, so local mutations diverge '
+                       'from platform results. Compute new tensors inside the hook instead.')
+        if _nest_fingerprint(prev_outputs) != self.last_outputs_fingerprint:
+            self._fail('the model outputs were mutated in place after the model call — the '
+                       'platform re-reads the stored step tensors, so local mutations diverge '
+                       'from platform results. Compute derived tensors inside the hook instead.')
+        if _nest_fingerprint(state) != self.last_state_fingerprint:
+            self._fail('the state object was mutated between hook calls — state may only change '
+                       'inside the hook body.')
+
+    def on_hook_return(self, next_inputs, state):
+        self.steps += 1
+        if self.steps > _MODEL_LOOP_MAX_STEPS:
+            self._fail(f'the chain exceeded {_MODEL_LOOP_MAX_STEPS} steps without the '
+                       f'autoregressive step returning None — the loop would never terminate on '
+                       f'the platform.')
+        self.last_state = state
+        self.last_state_fingerprint = _nest_fingerprint(state)
+        if next_inputs is None:
+            self.phase = 'ended'
+        else:
+            self.fed_inputs = next_inputs
+            self.fed_inputs_fingerprint = _nest_fingerprint(next_inputs)
+            self.phase = 'awaiting_model'
+
+    def mapping_hook_call(self, inputs_placeholder, state_placeholder):
+        if self.phase == 'ended':
+            self._fail('the chain already ended — once the autoregressive step returns None for '
+                       'next_inputs the loop must exit without further hook or model calls.')
+        if self.phase == 'awaiting_model':
+            self._fail('two autoregressive step calls in a row — the model must run exactly once '
+                       'between hook calls.')
+        self.last_state = state_placeholder
+        if self.phase == 'awaiting_first_hook':
+            self.fed_inputs = inputs_placeholder
+            self.phase = 'awaiting_model'
+            return inputs_placeholder, state_placeholder
+        self.phase = 'ended'
+        return None, state_placeholder
+
+    def on_model_call(self, fed):
+        if self.phase == 'awaiting_first_hook':
+            self._fail('the model was called before the autoregressive step produced the initial '
+                       'inputs — the loop must start with the initial hook call '
+                       '(prev_inputs=None, prev_outputs=None, state=None).')
+        if self.phase == 'awaiting_hook':
+            self._fail('the model was called twice in a row — the autoregressive step must run '
+                       'between model calls.')
+        if self.phase == 'ended':
+            self._fail('the chain already ended — once the autoregressive step returns None for '
+                       'next_inputs the loop must exit without further hook or model calls.')
+        if self.is_mapping:
+            return
+        fed_ids = {id(leaf) for leaf in _ndarray_leaves(fed)}
+        hook_ids = {id(value) for value in self.fed_inputs.values()}
+        if fed_ids != hook_ids:
+            self._fail('the model must be fed exactly the tensors the last hook call returned — '
+                       'any computation between the hook and the model is invisible to the '
+                       'platform. Move it into the hook.')
+        if _nest_fingerprint(self.fed_inputs) != self.fed_inputs_fingerprint:
+            self._fail('the model inputs were mutated in place after the hook returned them — '
+                       'the platform feeds the model the tensors exactly as the hook returned '
+                       'them. Compute new tensors inside the hook instead.')
+
+    def on_model_return(self, raw_outputs):
+        self.phase = 'awaiting_hook'
+        if self.is_mapping:
+            self.last_outputs = raw_outputs
+            return raw_outputs
+        outputs_list = raw_outputs if isinstance(raw_outputs, list) else [raw_outputs]
+        if len(outputs_list) != len(self.prediction_names):
+            self._fail(f'the model returned {len(outputs_list)} outputs but '
+                       f'{len(self.prediction_names)} prediction types are declared on '
+                       f'tensorleap_load_model — declare one prediction type per model output.')
+        named_outputs = {name: np.asarray(output)
+                         for name, output in zip(self.prediction_names, outputs_list)}
+        self.last_outputs = named_outputs
+        self.last_outputs_fingerprint = _nest_fingerprint(named_outputs)
+        return named_outputs
+
+
+class _ModelLoopModelProxy:
+    def __init__(self, model, context):
+        self._model = model
+        self._context = context
+
+    def __call__(self, inputs):
+        self._context.on_model_call(inputs)
+        return self._context.on_model_return(self._model(inputs))
+
+    def run(self, output_names, input_dict):
+        self._context.on_model_call(input_dict)
+        return self._context.on_model_return(self._model.run(output_names, input_dict))
+
+    def get_inputs(self):
+        return self._model.get_inputs()
+
+
+def tensorleap_autoregressive_step(latent_space_aggregation: str = 'last_step'):
+    """The feedback hook that drives an autoregressive chain.
+
+    Signature of the decorated function:
+        (sample_id, prev_inputs, prev_outputs, state, preprocess)
+            -> (Optional[Dict[str, np.ndarray]], state)
+
+    First call per chain: prev_inputs=None, prev_outputs=None, state=None — returns the initial
+    model inputs and the initial state. Later calls receive the previous step's model inputs and
+    outputs (unbatched, per-sample; outputs keyed by the declared prediction-type names) plus the
+    state returned by the previous call, and return the next step's full input dict together with
+    the updated state; returning (None, state) ends the chain — the terminating call still
+    returns state, so the final step participates in state aggregation. State is never fed to
+    the model; it may be any nest of dicts/lists holding numpy arrays, numbers, strings or bools,
+    and it rides the platform queue on every step — accumulate reductions, not raw per-step
+    tensors. The hook must be stateless and deterministically seeded (seed stochasticity from
+    sample_id) — the engine replays steps after crash recovery and relies on identical results.
+
+    latent_space_aggregation ('last_step' | 'mean') declares how the chain's latent-space
+    vectors are derived from its per-step forward passes. 'last_step' (default): every latent
+    space comes from the final step's forward pass, except input-kind latent spaces which come
+    from the first step — the original sample, before generated content dominates the model
+    inputs. 'mean': every latent space is the elementwise mean over all steps of the chain.
+    """
+    assert isinstance(latent_space_aggregation, str), \
+        ('tensorleap_autoregressive_step must be called with parentheses: '
+         '@tensorleap_autoregressive_step() or '
+         "@tensorleap_autoregressive_step(latent_space_aggregation='mean').")
+
+    def decorating_function(user_function: AutoregressiveStepCallableInterface):
+        # Contract: shapes and key set are fixed across all calls AND all chains (compiled graphs
+        # need static shapes), so one record shared across samples is the correct scope.
+        first_call_record: Dict[str, Any] = {'keys': None, 'shapes': None}
+
+        def _normalize_args(args, kwargs):
+            expected_names = ["sample_id", "prev_inputs", "prev_outputs", "state", "preprocess"]
+            normalized = []
+            for i, arg_name in enumerate(expected_names):
+                if i < len(args):
+                    normalized.append(args[i])
+                elif arg_name in kwargs:
+                    normalized.append(kwargs[arg_name])
+                else:
+                    raise AssertionError(
+                        f'{user_function.__name__} validation failed: missing required argument '
+                        f"'{arg_name}'. Expected arguments: {expected_names}.")
+            if len(args) + len(kwargs) != len(expected_names):
+                raise AssertionError(
+                    f'{user_function.__name__} validation failed: expected exactly '
+                    f'{len(expected_names)} arguments {expected_names}, got '
+                    f'{len(args) + len(kwargs)}.')
+            sample_id, prev_inputs, prev_outputs, state, preprocess_response = normalized
+            assert isinstance(sample_id, (int, str)), \
+                (f'{user_function.__name__}() validation failed: '
+                 f'sample_id must be an int or str. Got {type(sample_id).__name__}.')
+            assert prev_inputs is None or isinstance(prev_inputs, dict), \
+                (f'{user_function.__name__}() validation failed: '
+                 f'prev_inputs must be a dict or None. Got {type(prev_inputs).__name__}.')
+            assert prev_outputs is None or isinstance(prev_outputs, dict), \
+                (f'{user_function.__name__}() validation failed: '
+                 f'prev_outputs must be a dict or None. Got {type(prev_outputs).__name__}.')
+            assert isinstance(preprocess_response, PreprocessResponse), \
+                (f'{user_function.__name__}() validation failed: '
+                 f'preprocess must be a PreprocessResponse. Got '
+                 f'{type(preprocess_response).__name__}.')
+            return normalized
+
+        def _validate_input_args(sample_id, prev_inputs, prev_outputs, state, preprocess_response):
+            assert (prev_inputs is None) == (prev_outputs is None), \
+                (f'{user_function.__name__}() validation failed: '
+                 f'prev_inputs and prev_outputs must both be None (first call) or both be dicts '
+                 f'(subsequent calls). Got prev_inputs={type(prev_inputs).__name__}, '
+                 f'prev_outputs={type(prev_outputs).__name__}.')
+            assert not (prev_inputs is None and state is not None), \
+                (f'{user_function.__name__}() validation failed: '
+                 f'the first call (prev_inputs=None, prev_outputs=None) must pass state=None. '
+                 f'Got state of type {type(state).__name__}.')
+            assert type(sample_id) == preprocess_response.sample_id_type, \
+                (f'{user_function.__name__}() validation failed: '
+                 f'Argument sample_id should be as the same type as defined in the preprocess response '
+                 f'{preprocess_response.sample_id_type}. Got {type(sample_id)}.')
+
+        def _validate_result(result, is_first_call):
+            assert isinstance(result, tuple) and len(result) == 2, \
+                (f'{user_function.__name__}() validation failed: '
+                 f'the hook must return a (next_inputs, state) tuple — next_inputs is a dict of '
+                 f'named model input tensors, or None to end the chain; state rides to the next '
+                 f'call. Got {type(result).__name__}.')
+            next_inputs, state = result
+            if next_inputs is None:
+                assert not is_first_call, \
+                    (f'{user_function.__name__}() validation failed: '
+                     f'the first call (prev_inputs=None, prev_outputs=None, state=None) must return '
+                     f'the initial model inputs dict — returning None on the first call would '
+                     f'produce an empty chain.')
+                return
+            assert isinstance(next_inputs, dict), \
+                (f'{user_function.__name__}() validation failed: '
+                 f'next_inputs must be a dict of named model input tensors, or None to end the '
+                 f'chain. Got {type(next_inputs)}.')
+            assert next_inputs, \
+                (f'{user_function.__name__}() validation failed: '
+                 f'next_inputs is an empty dict — return the full model input dict, or None to '
+                 f'end the chain.')
+            for key, value in next_inputs.items():
+                assert isinstance(key, str), \
+                    (f'{user_function.__name__}() validation failed: '
+                     f'all keys in next_inputs must be strings. Got {type(key)}.')
+                assert not key.startswith('_'), \
+                    (f'{user_function.__name__}() validation failed: '
+                     f'next_inputs key "{key}" starts with "_" — underscore passthrough keys were '
+                     f'replaced by the explicit state channel. Move passthrough values into the '
+                     f'returned state.')
+                assert isinstance(value, np.ndarray), \
+                    (f'{user_function.__name__}() validation failed: '
+                     f'value for model input "{key}" must be a numpy array. Got {type(value)}.')
+
+        def _validate_consistency(next_inputs):
+            # Rules 2-3 of the autoregressive contract: identical key set and identical shapes on
+            # every call — compiled graphs need static shapes.
+            if next_inputs is None:
+                return
+            keys = set(next_inputs.keys())
+            shapes = {key: tuple(value.shape) for key, value in next_inputs.items()}
+            if first_call_record['keys'] is None:
+                first_call_record['keys'] = keys
+                first_call_record['shapes'] = shapes
+                return
+            assert keys == first_call_record['keys'], \
+                (f'{user_function.__name__}() validation failed: '
+                 f'every call must return the same key set as the first call. '
+                 f'First call keys: {sorted(first_call_record["keys"])}, this call: {sorted(keys)}.')
+            for key, shape in shapes.items():
+                assert shape == first_call_record['shapes'][key], \
+                    (f'{user_function.__name__}() validation failed: '
+                     f'tensor shapes must be identical on every call (fixed-length padding + mask for '
+                     f'growing sequences). Input "{key}" was {first_call_record["shapes"][key]} on the '
+                     f'first call and {shape} now.')
+
+        def _assert_deterministic(first_result, second_result):
+            error_message = (
+                f'{user_function.__name__}() validation failed: '
+                f'the hook is not deterministic — two calls with identical arguments returned different '
+                f'outputs. Seed any stochasticity (e.g. initial noise) from sample_id: the engine '
+                f'replays steps after crash recovery and relies on identical results.')
+            assert isinstance(first_result, tuple) and isinstance(second_result, tuple) and \
+                len(first_result) == 2 and len(second_result) == 2, error_message
+            first_inputs, first_state = first_result
+            second_inputs, second_state = second_result
+            if first_inputs is None or second_inputs is None:
+                assert first_inputs is None and second_inputs is None, error_message
+            else:
+                assert set(first_inputs.keys()) == set(second_inputs.keys()), error_message
+                for key, value in first_inputs.items():
+                    if not isinstance(value, np.ndarray):
+                        continue
+                    assert np.array_equal(value, second_inputs[key]), error_message
+            assert _nest_equal(first_state, second_state), error_message
+
+        def _squeeze_test_batch_dim(tensors_dict):
+            # Inside the integration test the model works on batch-1 tensors (this wrapper adds the
+            # batch axis to returned inputs, and model outputs come back batched). The user's hook
+            # body must see the same unbatched per-sample tensors it will see at engine runtime, so
+            # strip the leading batch-1 axis before calling it.
+            if tensors_dict is None:
+                return None
+            squeezed = {}
+            for key, value in tensors_dict.items():
+                if isinstance(value, np.ndarray) and value.ndim > 0 and value.shape[0] == 1:
+                    squeezed[key] = value[0]
+                else:
+                    squeezed[key] = value
+            return squeezed
+
+        def inner_without_validate(sample_id, prev_inputs, prev_outputs, state, preprocess_response):
+            global _called_from_inside_tl_decorator
+            _called_from_inside_tl_decorator += 1
+            try:
+                result = user_function(sample_id, prev_inputs, prev_outputs, state,
+                                       preprocess_response)
+            finally:
+                _called_from_inside_tl_decorator -= 1
+            return result
+
+        leap_binder.set_autoregressive_step(inner_without_validate,
+                                            latent_space_aggregation=latent_space_aggregation)
+
+        def inner(*args, **kwargs):
+            if not _call_from_tl_platform:
+                set_current('tensorleap_autoregressive_step')
+            sample_id, prev_inputs, prev_outputs, state, preprocess_response = \
+                _normalize_args(args, kwargs)
+
+            is_top_level_in_test = (_called_from_inside_tl_decorator == 0
+                                    and _called_from_inside_tl_integration_test_decorator)
+            loop_context = _active_model_loop
+            if is_top_level_in_test:
+                if loop_context is None:
+                    raise LeapValidationError(
+                        f'{user_function.__name__}() validation failed: inside the integration '
+                        f'test the autoregressive step may only be called from within the '
+                        f'tensorleap_model_loop function — including the initial call '
+                        f'(prev_inputs=None). The platform drives the chain through the model '
+                        f'loop only, so a hook call outside it would not run on the platform.')
+                loop_context.on_hook_call(sample_id, prev_inputs, prev_outputs, state,
+                                          preprocess_response)
+                prev_inputs = _squeeze_test_batch_dim(prev_inputs)
+                prev_outputs = _squeeze_test_batch_dim(prev_outputs)
+
+            _validate_input_args(sample_id, prev_inputs, prev_outputs, state, preprocess_response)
+
+            # The deep copy keeps the determinism double-call independent: a hook that
+            # accumulates into state in place would otherwise see its own first-call mutations.
+            state_for_second_call = copy.deepcopy(state) if is_top_level_in_test else None
+            result = inner_without_validate(sample_id, prev_inputs, prev_outputs, state,
+                                            preprocess_response)
+
+            if is_top_level_in_test:
+                # Rule 7: determinism double-call. An unseeded sampler fails here, on the user's
+                # machine, instead of silently forking a chain on engine crash-replay.
+                second_result = inner_without_validate(sample_id, prev_inputs, prev_outputs,
+                                                       state_for_second_call, preprocess_response)
+                _assert_deterministic(result, second_result)
+
+            _validate_result(result, is_first_call=prev_inputs is None)
+            next_inputs, new_state = result
+            _validate_consistency(next_inputs)
+            _validate_state_types(new_state)
+
+            if is_top_level_in_test:
+                if next_inputs is not None:
+                    next_inputs = {key: np.expand_dims(value, axis=0)
+                                   for key, value in next_inputs.items()}
+                loop_context.on_hook_return(next_inputs, new_state)
+
+            if not _call_from_tl_platform:
+                update_env_params_func("tensorleap_autoregressive_step", "v")
+            return next_inputs, new_state
+
+        class _MappingInputsPlaceholder:
+            """Mapping-mode stand-in for the hook's returned inputs dict.
+
+            Indexing it by a model-input name yields a placeholder whose node_mapping the model
+            placeholder tags Input0/Input1/... — exactly how input encoders wire the graph. One
+            NodeConnection is recorded per key, once.
+            """
+            def __init__(self):
+                self._items: Dict[str, Any] = {}
+
+            def __getitem__(self, key):
+                assert isinstance(key, str), \
+                    f'Expected a string model input name, got {type(key)} instead.'
+                if key not in self._items:
+                    class TempMapping:
+                        pass
+                    ret = TempMapping()
+                    ret.node_mapping = NodeMapping(key, NodeMappingType.Input)
+                    leap_binder.mapping_connections.append(NodeConnection(ret.node_mapping, None))
+                    self._items[key] = ret
+                return self._items[key]
+
+            def get(self, key, default=None):
+                # Deliberate alias of [] — the fixed-key-set contract makes a defaulted lookup
+                # meaningless in mapping mode, and every accessed key must wire into the graph.
+                return self[key]
+
+            def _unsupported_iteration(self):
+                raise Exception(
+                    'Iterating over the autoregressive hook result is not supported inside the '
+                    'integration test. Pass the model its inputs with explicit keys, e.g. '
+                    "model.run(None, {'x': inputs['x'], 't': inputs['t']}).")
+
+            def items(self):
+                self._unsupported_iteration()
+
+            def keys(self):
+                self._unsupported_iteration()
+
+            def values(self):
+                self._unsupported_iteration()
+
+        mapping_placeholder = _MappingInputsPlaceholder()
+        mapping_state_placeholder = _MappingStatePlaceholder()
+
+        def mapping_inner(*args, **kwargs):
+            loop_context = _active_model_loop
+            if loop_context is None:
+                raise LeapValidationError(
+                    f'{user_function.__name__}() validation failed: inside the integration test '
+                    f'the autoregressive step may only be called from within the '
+                    f'tensorleap_model_loop function — including the initial call '
+                    f'(prev_inputs=None).')
+            return loop_context.mapping_hook_call(mapping_placeholder, mapping_state_placeholder)
+
+        def final_inner(*args, **kwargs):
+            if os.environ.get(mapping_runtime_mode_env_var_mame):
+                return mapping_inner(*args, **kwargs)
+            else:
+                return inner(*args, **kwargs)
+
+        if not _call_from_tl_platform:
+            register_decorator_row("tensorleap_autoregressive_step")
+
+        return final_inner
+
+    return decorating_function
+
+
+def tensorleap_model_loop():
+    """The user-authored autoregressive rollout loop.
+
+    Signature of the decorated function:
+        (model, sample_id, preprocess) -> (final_inputs, final_outputs, final_state)
+
+    The function drives the chain locally, exactly as the platform will: it makes the initial
+    hook call (prev_inputs=None, prev_outputs=None, state=None) inside the loop, alternates
+    model forwards with hook calls, feeds the model exactly the tensors the hook returned, and
+    exits when the hook returns None for next_inputs — returning the last inputs dict fed to the
+    model, the last model outputs (a dict keyed by the declared prediction-type names) and the
+    final state. The loop body itself never runs on the platform — the engine drives the chain —
+    so the decorator validates, while the loop runs, that its behavior matches the engine loop;
+    any computation in the loop body outside the hook and the model is invisible to the platform
+    and fails validation. Canonical shape:
+
+        @tensorleap_model_loop()
+        def model_loop(model, sample_id, preprocess):
+            inputs, state = generate_next_inputs(sample_id, None, None, None, preprocess)
+            while True:
+                outputs = model(inputs)
+                next_inputs, state = generate_next_inputs(sample_id, inputs, outputs, state,
+                                                          preprocess)
+                if next_inputs is None:
+                    return inputs, outputs, state
+                inputs = next_inputs
+    """
+
+    def decorating_function(user_function):
+
+        def _normalize_args(args, kwargs):
+            expected_names = ["model", "sample_id", "preprocess"]
+            normalized = []
+            for i, arg_name in enumerate(expected_names):
+                if i < len(args):
+                    normalized.append(args[i])
+                elif arg_name in kwargs:
+                    normalized.append(kwargs[arg_name])
+                else:
+                    raise AssertionError(
+                        f'{user_function.__name__} validation failed: missing required argument '
+                        f"'{arg_name}'. Expected arguments: {expected_names}.")
+            if len(args) + len(kwargs) != len(expected_names):
+                raise AssertionError(
+                    f'{user_function.__name__} validation failed: expected exactly '
+                    f'{len(expected_names)} arguments {expected_names}, got '
+                    f'{len(args) + len(kwargs)}.')
+            return normalized
+
+        def _fail(message):
+            raise LeapValidationError(f'tensorleap_model_loop validation failed: {message}')
+
+        def _run_with_context(model, sample_id, preprocess_response, is_mapping):
+            global _active_model_loop, _model_loop_run_in_current_test
+            if _active_model_loop is not None:
+                _fail('model loops cannot nest or run concurrently.')
+            if _model_loop_run_in_current_test:
+                _fail('only one model loop run is allowed per integration test.')
+            if model is None or model is not _last_loaded_model:
+                _fail('the model argument must be the object returned by the '
+                      'tensorleap_load_model function, unchanged.')
+            if leap_binder.setup_container.autoregressive_step is None:
+                _fail('tensorleap_model_loop requires a tensorleap_autoregressive_step hook.')
+            prediction_names = [prediction_type.name for prediction_type
+                                in leap_binder.setup_container.prediction_types]
+            if not prediction_names:
+                _fail('tensorleap_model_loop requires prediction types declared on '
+                      'tensorleap_load_model — the loop keys the model outputs by their names.')
+            context = _ModelLoopContext(sample_id, preprocess_response, prediction_names,
+                                        is_mapping)
+            _active_model_loop = context
+            try:
+                result = user_function(_ModelLoopModelProxy(model, context), sample_id,
+                                       preprocess_response)
+            finally:
+                _active_model_loop = None
+            _model_loop_run_in_current_test = True
+            if context.phase == 'awaiting_first_hook':
+                _fail('the loop never called the autoregressive step — it must start with the '
+                      'initial hook call (prev_inputs=None, prev_outputs=None, state=None).')
+            if context.phase != 'ended':
+                _fail('the loop returned before the autoregressive step signaled the chain end — '
+                      'it must run until the hook returns None for next_inputs and then return '
+                      '(final_inputs, final_outputs, final_state).')
+            if not isinstance(result, tuple) or len(result) != 3:
+                _fail('the loop must return (final_inputs, final_outputs, final_state). Got '
+                      f'{type(result).__name__}.')
+            final_inputs, final_outputs, final_state = result
+            if final_inputs is not context.fed_inputs:
+                _fail('final_inputs must be exactly the last inputs dict fed to the model — the '
+                      'one the hook returned before the terminating call.')
+            if final_outputs is not context.last_outputs:
+                _fail('final_outputs must be exactly the outputs dict of the last model call.')
+            if final_state is not context.last_state:
+                _fail('final_state must be exactly the state returned by the terminating hook '
+                      'call.')
+            return context, final_inputs, final_outputs, final_state
+
+        def inner(*args, **kwargs):
+            if not _call_from_tl_platform:
+                set_current('tensorleap_model_loop')
+            model, sample_id, preprocess_response = _normalize_args(args, kwargs)
+            in_test = _called_from_inside_tl_integration_test_decorator
+            context, final_inputs, final_outputs, final_state = _run_with_context(
+                model, sample_id, preprocess_response, is_mapping=False)
+            if _nest_fingerprint(final_inputs) != context.fed_inputs_fingerprint:
+                _fail('the final inputs were mutated in place after the last model call.')
+            if _nest_fingerprint(final_outputs) != context.last_outputs_fingerprint:
+                _fail('the final outputs were mutated in place after the last model call.')
+            if _nest_fingerprint(final_state) != context.last_state_fingerprint:
+                _fail('the final state was mutated in place after the terminating hook call.')
+            if not in_test:
+                if not _call_from_tl_platform:
+                    update_env_params_func("tensorleap_model_loop", "v")
+                return final_inputs, final_outputs, final_state
+            # Inside the integration test the chain tensors carry a batch-1 axis; the platform
+            # feeds autoregressive metrics per-chain unbatched dicts, so strip it here.
+            squeezed_inputs = _squeeze_leading_batch_axis(final_inputs)
+            squeezed_outputs = _squeeze_leading_batch_axis(final_outputs)
+            _register_chain_artifact(squeezed_inputs, 'final_inputs')
+            _register_chain_artifact(squeezed_outputs, 'final_outputs')
+            _register_chain_artifact(final_state, 'final_state')
+            if not _call_from_tl_platform:
+                update_env_params_func("tensorleap_model_loop", "v")
+            return squeezed_inputs, squeezed_outputs, final_state
+
+        def mapping_inner(*args, **kwargs):
+            model, sample_id, preprocess_response = _normalize_args(args, kwargs)
+            context, final_inputs, final_outputs, final_state = _run_with_context(
+                model, sample_id, preprocess_response, is_mapping=True)
+            _register_chain_artifact(final_inputs, 'final_inputs')
+            _register_chain_artifact(final_outputs, 'final_outputs')
+            _register_chain_artifact(final_state, 'final_state')
+            return final_inputs, final_outputs, final_state
+
+        def final_inner(*args, **kwargs):
+            if os.environ.get(mapping_runtime_mode_env_var_mame):
+                return mapping_inner(*args, **kwargs)
+            else:
+                return inner(*args, **kwargs)
+
+        if not _call_from_tl_platform:
+            register_decorator_row("tensorleap_model_loop")
+
+        return final_inner
+
+    return decorating_function
+
+
+_IMPLICIT_ARG_TO_ARTIFACT_KIND = {'inputs': 'final_inputs', 'outputs': 'final_outputs',
+                                  'state': 'final_state'}
+_ARTIFACT_KIND_DESCRIPTIONS = {
+    'final_inputs': 'final inputs dict returned by the model loop',
+    'final_outputs': 'final outputs dict returned by the model loop',
+    'final_state': 'final state returned by the model loop',
+    'ground_truth': 'direct return value of a ground-truth encoder',
+}
+_AUTOREGRESSIVE_SCALAR_TYPES = (int, float, np.floating, np.integer, np.bool_)
+
+
+def _normalize_autoregressive_call_args(func_name, full_arg_names, args, kwargs):
+    if len(args) + len(kwargs) != len(full_arg_names):
+        raise AssertionError(
+            f'{func_name} validation failed: expected exactly {len(full_arg_names)} arguments '
+            f'{full_arg_names}, got {len(args) + len(kwargs)}.')
+    normalized = {}
+    for i, arg_name in enumerate(full_arg_names):
+        if i < len(args):
+            normalized[arg_name] = args[i]
+        elif arg_name in kwargs:
+            normalized[arg_name] = kwargs[arg_name]
+        else:
+            raise AssertionError(
+                f'{func_name} validation failed: missing required argument \'{arg_name}\'. '
+                f'Expected arguments: {full_arg_names}.')
+    return normalized
+
+
+def _validate_autoregressive_arg_provenance(func_name, normalized):
+    for arg_name, value in normalized.items():
+        expected_kind = _IMPLICIT_ARG_TO_ARTIFACT_KIND.get(arg_name, 'ground_truth')
+        entry = _lookup_chain_artifact(value)
+        if entry is None:
+            raise LeapValidationError(
+                f'{func_name} validation failed: argument "{arg_name}" must be the '
+                f'{_ARTIFACT_KIND_DESCRIPTIONS[expected_kind]}, passed whole and unmodified — '
+                f'the platform feeds it from the finished chain, so anything else would not '
+                f'exist at platform runtime. Slice or derive inside the function body instead.')
+        kind, fingerprint, _ = entry
+        if kind != expected_kind:
+            raise LeapValidationError(
+                f'{func_name} validation failed: argument "{arg_name}" received the '
+                f'{_ARTIFACT_KIND_DESCRIPTIONS[kind]} but must be the '
+                f'{_ARTIFACT_KIND_DESCRIPTIONS[expected_kind]}.')
+        if fingerprint != _nest_fingerprint(value):
+            raise LeapValidationError(
+                f'{func_name} validation failed: argument "{arg_name}" was mutated in place '
+                f'after it was produced — the platform feeds the original values, so local '
+                f'mutations diverge from platform results.')
+
+
+def _squeeze_wired_autoregressive_args(normalized):
+    squeezed = {}
+    for arg_name, value in normalized.items():
+        if arg_name not in _IMPLICIT_ARG_TO_ARTIFACT_KIND and isinstance(value, np.ndarray) \
+                and value.ndim > 0 and value.shape[0] == 1:
+            squeezed[arg_name] = value[0]
+        else:
+            squeezed[arg_name] = value
+    return squeezed
+
+
+def _validate_autoregressive_scalar_result(func_name, result, allow_dict):
+    def _is_scalar(value):
+        return isinstance(value, _AUTOREGRESSIVE_SCALAR_TYPES) or \
+            (isinstance(value, np.ndarray) and value.ndim == 0)
+
+    if _is_scalar(result):
+        return
+    if allow_dict and isinstance(result, dict) and result and \
+            all(isinstance(key, str) for key in result) and \
+            all(_is_scalar(value) for value in result.values()):
+        return
+    expected = 'a scalar number' + (' or a dict of named scalar numbers' if allow_dict else '')
+    raise AssertionError(
+        f'{func_name} validation failed: autoregressive results are per-chain — the function '
+        f'must return {expected} for the single chain it received. Got {type(result).__name__}.')
+
+
+def _autoregressive_mapping_call(mapping_inner_func, args, kwargs, node_mapping_type):
+    user_unique_name = mapping_inner_func.name
+    if 'user_unique_name' in kwargs:
+        kwargs = dict(kwargs)
+        user_unique_name = kwargs.pop('user_unique_name')
+    if user_unique_name in mapping_inner_func.name_to_unique_name[mapping_inner_func.name]:
+        user_unique_name = \
+            f'{user_unique_name}_{len(mapping_inner_func.name_to_unique_name[mapping_inner_func.name])}'
+    mapping_inner_func.name_to_unique_name[mapping_inner_func.name].add(user_unique_name)
+
+    normalized = _normalize_autoregressive_call_args(mapping_inner_func.name,
+                                                     mapping_inner_func.full_arg_names, args,
+                                                     kwargs)
+    for arg_name, value in normalized.items():
+        if arg_name in _IMPLICIT_ARG_TO_ARTIFACT_KIND:
+            entry = _lookup_chain_artifact(value)
+            expected_kind = _IMPLICIT_ARG_TO_ARTIFACT_KIND[arg_name]
+            if entry is None or entry[0] != expected_kind:
+                raise LeapValidationError(
+                    f'{mapping_inner_func.name} validation failed: argument "{arg_name}" must '
+                    f'be the {_ARTIFACT_KIND_DESCRIPTIONS[expected_kind]}, passed whole.')
+        elif not hasattr(value, 'node_mapping'):
+            raise LeapValidationError(
+                f'{mapping_inner_func.name} validation failed: argument "{arg_name}" must be '
+                f'the direct return value of a ground-truth encoder.')
+    ordered_connections = [normalized[arg_name] for arg_name in mapping_inner_func.arg_names]
+    _add_mapping_connection(user_unique_name, ordered_connections, mapping_inner_func.arg_names,
+                            mapping_inner_func.name, node_mapping_type)
+    return None
+
+
+def _make_autoregressive_decorator_inner(user_function, name, decorator_env_name, handler_data,
+                                         node_mapping_type, validate_result):
+    full_arg_names = inspect.getfullargspec(user_function)[0]
+
+    def inner_without_validate(*args, **kwargs):
+        global _called_from_inside_tl_decorator
+        _called_from_inside_tl_decorator += 1
+        try:
+            result = user_function(*args, **kwargs)
+        finally:
+            _called_from_inside_tl_decorator -= 1
+        return result
+
+    def inner(*args, **kwargs):
+        if not _call_from_tl_platform:
+            set_current(decorator_env_name)
+        normalized = _normalize_autoregressive_call_args(user_function.__name__, full_arg_names,
+                                                         args, kwargs)
+        is_top_level_in_test = (_called_from_inside_tl_decorator == 0
+                                and _called_from_inside_tl_integration_test_decorator)
+        if is_top_level_in_test:
+            _validate_autoregressive_arg_provenance(user_function.__name__, normalized)
+            normalized = _squeeze_wired_autoregressive_args(normalized)
+        result = inner_without_validate(**normalized)
+        validate_result(user_function.__name__, result)
+        if not _call_from_tl_platform:
+            update_env_params_func(decorator_env_name, "v")
+        return result
+
+    def mapping_inner(*args, **kwargs):
+        return _autoregressive_mapping_call(mapping_inner, args, kwargs, node_mapping_type)
+
+    mapping_inner.name = name
+    mapping_inner.full_arg_names = full_arg_names
+    mapping_inner.arg_names = handler_data.arg_names
+    mapping_inner.name_to_unique_name = defaultdict(set)
+
+    def final_inner(*args, **kwargs):
+        if os.environ.get(mapping_runtime_mode_env_var_mame):
+            return mapping_inner(*args, **kwargs)
+        else:
+            return inner(*args, **kwargs)
+
+    if not _call_from_tl_platform:
+        register_decorator_row(decorator_env_name)
+
+    return final_inner
+
+
+def tensorleap_autoregressive_metric(name: str,
+                                     direction: Union[MetricDirection, Dict[str, MetricDirection]]
+                                     = MetricDirection.Downward,
+                                     compute_insights: Optional[Union[bool, Dict[str, bool]]] = None):
+    """A per-chain metric over a finished autoregressive chain.
+
+    The reserved arguments inputs / outputs / state are fed implicitly by the platform from the
+    chain's final step (unbatched dicts; outputs keyed by the declared prediction-type names;
+    state as the terminating hook call returned it). Declare whichever of them the metric needs.
+    Every other argument is wired to a ground-truth encoder through the integration test. Runs
+    once per chain and must return a scalar number (or a dict of named scalar numbers).
+    """
+
+    def decorating_function(user_function):
+        leap_binder.add_autoregressive_metric(user_function, name, direction, compute_insights)
+        handler_data = \
+            leap_binder.setup_container.autoregressive_metrics[-1].metric_handler_data
+
+        def validate_result(func_name, result):
+            _validate_autoregressive_scalar_result(func_name, result, allow_dict=True)
+
+        return _make_autoregressive_decorator_inner(
+            user_function, name, 'tensorleap_autoregressive_metric', handler_data,
+            NodeMappingType.Metric, validate_result)
+
+    return decorating_function
+
+
+def tensorleap_autoregressive_loss(name: str):
+    """A per-chain loss over a finished autoregressive chain.
+
+    Same argument contract as tensorleap_autoregressive_metric; must return a single scalar
+    number for the chain it received.
+    """
+
+    def decorating_function(user_function):
+        leap_binder.add_autoregressive_loss(user_function, name)
+        handler_data = \
+            leap_binder.setup_container.autoregressive_losses[-1].custom_loss_handler_data
+
+        def validate_result(func_name, result):
+            _validate_autoregressive_scalar_result(func_name, result, allow_dict=False)
+
+        return _make_autoregressive_decorator_inner(
+            user_function, name, 'tensorleap_autoregressive_loss', handler_data,
+            NodeMappingType.CustomLoss, validate_result)
+
+    return decorating_function
+
+
+def tensorleap_autoregressive_visualizer(name: str, visualizer_type: LeapDataType):
+    """A per-chain visualizer over a finished autoregressive chain.
+
+    Same argument contract as tensorleap_autoregressive_metric; must return a LeapData object
+    matching visualizer_type.
+    """
+
+    def decorating_function(user_function):
+        assert isinstance(visualizer_type, LeapDataType), \
+            (f'{user_function.__name__} validation failed: visualizer_type should be of type '
+             f'{LeapDataType.__name__} but got {type(visualizer_type)}')
+        leap_binder.add_autoregressive_visualizer(user_function, name, visualizer_type)
+        handler_data = \
+            leap_binder.setup_container.autoregressive_visualizers[-1].visualizer_handler_data
+
+        def validate_result(func_name, result):
+            expected_class = map_leap_data_type_to_visualizer_class[visualizer_type.value]
+            assert isinstance(result, expected_class), \
+                (f'{func_name} validation failed: the return type should be '
+                 f'{expected_class.__name__} (visualizer_type={visualizer_type}). '
+                 f'Got {type(result).__name__}.')
+
+        return _make_autoregressive_decorator_inner(
+            user_function, name, 'tensorleap_autoregressive_visualizer', handler_data,
+            NodeMappingType.Visualizer, validate_result)
+
+    return decorating_function
+
+
 def tensorleap_preprocess():
     def decorating_function(user_function: Callable[[], List[PreprocessResponse]]):
         leap_binder.set_preprocess(user_function)
@@ -1732,15 +2620,6 @@ def tensorleap_preprocess():
             assert len(set(result)) == len(result), \
                 (f'{user_function.__name__}() validation failed: '
                  f'The return list should not contain duplicate PreprocessResponse objects.')
-            # All-or-none: a dataset is either all grouped (nested sample_ids) or all flat.
-            if len({response.is_grouped for response in result}) != 1:
-                shapes = ', '.join(
-                    f"#{i}={'grouped' if response.is_grouped else 'flat'}"
-                    for i, response in enumerate(result))
-                raise AssertionError(
-                    f'{user_function.__name__}() validation failed: '
-                    f'All PreprocessResponses must be either all grouped (nested sample_ids) '
-                    f'or all flat, but got a mix: {shapes}.')
 
         def inner(*args, **kwargs):
             if not _call_from_tl_platform:
@@ -2078,10 +2957,14 @@ def tensorleap_input_encoder(name: str, channel_dim=_UNSET, model_input_index=No
                 f"channel axis in the batched tensor (batch = axis 0), so 0 is invalid. Use 1 "
                 f"for channel-first (NCHW) and -1 for channel-last (NHWC).")
 
-        def _validate_input_args(sample_id: Union[int, str, list], preprocess_response: PreprocessResponse):
-            _validate_id_or_group(sample_id, preprocess_response, user_function.__name__)
+        def _validate_input_args(sample_id: Union[int, str], preprocess_response: PreprocessResponse):
+            assert type(sample_id) == preprocess_response.sample_id_type, \
+                (f'{user_function.__name__}() validation failed: '
+                 f'Argument sample_id should be as the same type as defined in the preprocess response '
+                 f'{preprocess_response.sample_id_type}. Got {type(sample_id)}.')
 
-        def _validate_single(result):
+        def _validate_result(result):
+            validate_output_structure(result, func_name=user_function.__name__, expected_type_name="np.ndarray")
             assert isinstance(result, np.ndarray), \
                 (f'{user_function.__name__}() validation failed: '
                  f'Unsupported return type. Should be a numpy array. Got {type(result)}.')
@@ -2091,13 +2974,6 @@ def tensorleap_input_encoder(name: str, channel_dim=_UNSET, model_input_index=No
             assert channel_dim - 1 <= len(result.shape), (f'{user_function.__name__}() validation failed: '
                                                           f'The channel_dim ({channel_dim}) should be <= to the rank of the resulting input rank ({len(result.shape)}).')
 
-        def _validate_result(result, grouped=False, group_size=None):
-            if not grouped:
-                validate_output_structure(result, func_name=user_function.__name__, expected_type_name="np.ndarray")
-                _validate_single(result)
-                return
-            _validate_grouped_result(result, group_size, user_function.__name__, _validate_single)
-
         def inner_without_validate(sample_id, preprocess_response):
             global _called_from_inside_tl_decorator
             _called_from_inside_tl_decorator += 1
@@ -2114,18 +2990,16 @@ def tensorleap_input_encoder(name: str, channel_dim=_UNSET, model_input_index=No
         def inner(*args, **kwargs):
             if not _call_from_tl_platform:
                 set_current("tensorleap_input_encoder")
-            validate_args_structure(*args, types_order=[Union[int, str, list], PreprocessResponse],
+            validate_args_structure(*args, types_order=[Union[int, str], PreprocessResponse],
                                     func_name=user_function.__name__, expected_names=["idx", "preprocess"], **kwargs)
             sample_id, preprocess_response = args if len(args) != 0 else kwargs.values()
             _validate_input_args(sample_id, preprocess_response)
 
-            grouped = isinstance(sample_id, list)
-            group_size = len(sample_id) if grouped else None
             result = inner_without_validate(sample_id, preprocess_response)
 
-            _validate_result(result, grouped, group_size)
+            _validate_result(result)
 
-            if not grouped and _called_from_inside_tl_decorator == 0 and _called_from_inside_tl_integration_test_decorator:
+            if _called_from_inside_tl_decorator == 0 and _called_from_inside_tl_integration_test_decorator:
                 batch_warning(result, user_function.__name__)
                 result = np.expand_dims(result, axis=0)
                 # Emit integration test event once per test
@@ -2148,15 +3022,8 @@ def tensorleap_input_encoder(name: str, channel_dim=_UNSET, model_input_index=No
         inner.node_mapping = NodeMapping(name, node_mapping_type)
 
         def mapping_inner(*args, **kwargs):
-            if _mapping_dataset_is_grouped:
-                class TempMapping:
-                    def __getitem__(self, key):
-                        # Indexing a grouped input group (input_list[i]) selects a sample from the
-                        # same input source; the model wiring is identical to the un-grouped input.
-                        return self
-            else:
-                class TempMapping:
-                    pass
+            class TempMapping:
+                pass
 
             ret = TempMapping()
             ret.node_mapping = mapping_inner.node_mapping
@@ -2186,10 +3053,15 @@ def tensorleap_gt_encoder(name: str):
                 raise Exception(f'GT with name {name} already exists. '
                                 f'Please choose another')
 
-        def _validate_input_args(sample_id: Union[int, str, list], preprocess_response: PreprocessResponse):
-            _validate_id_or_group(sample_id, preprocess_response, user_function.__name__)
+        def _validate_input_args(sample_id: Union[int, str], preprocess_response: PreprocessResponse):
+            assert type(sample_id) == preprocess_response.sample_id_type, \
+                (f'{user_function.__name__}() validation failed: '
+                 f'Argument sample_id should be as the same type as defined in the preprocess response '
+                 f'{preprocess_response.sample_id_type}. Got {type(sample_id)}.')
 
-        def _validate_single(result):
+        def _validate_result(result):
+            validate_output_structure(result, func_name=user_function.__name__, expected_type_name="np.ndarray",
+                                      gt_flag=True)
             assert isinstance(result, np.ndarray), \
                 (f'{user_function.__name__}() validation failed: '
                  f'Unsupported return type. Should be a numpy array. Got {type(result)}.')
@@ -2197,14 +3069,6 @@ def tensorleap_gt_encoder(name: str):
                 (f'{user_function.__name__}() validation failed: '
                  f'The return type should be a numpy array of type float32. Got {result.dtype}.')
 
-        def _validate_result(result, grouped=False, group_size=None):
-            if not grouped:
-                validate_output_structure(result, func_name=user_function.__name__, expected_type_name="np.ndarray",
-                                          gt_flag=True)
-                _validate_single(result)
-                return
-            _validate_grouped_result(result, group_size, user_function.__name__, _validate_single)
-
         def inner_without_validate(sample_id, preprocess_response):
             global _called_from_inside_tl_decorator
             _called_from_inside_tl_decorator += 1
@@ -2221,20 +3085,19 @@ def tensorleap_gt_encoder(name: str):
         def inner(*args, **kwargs):
             if not _call_from_tl_platform:
                 set_current("tensorleap_gt_encoder")
-            validate_args_structure(*args, types_order=[Union[int, str, list], PreprocessResponse],
+            validate_args_structure(*args, types_order=[Union[int, str], PreprocessResponse],
                                     func_name=user_function.__name__, expected_names=["idx", "preprocess"], **kwargs)
             sample_id, preprocess_response = args
             _validate_input_args(sample_id, preprocess_response)
 
-            grouped = isinstance(sample_id, list)
-            group_size = len(sample_id) if grouped else None
             result = inner_without_validate(sample_id, preprocess_response)
 
-            _validate_result(result, grouped, group_size)
+            _validate_result(result)
 
-            if not grouped and _called_from_inside_tl_decorator == 0 and _called_from_inside_tl_integration_test_decorator:
+            if _called_from_inside_tl_decorator == 0 and _called_from_inside_tl_integration_test_decorator:
                 batch_warning(result, user_function.__name__)
                 result = np.expand_dims(result, axis=0)
+                _register_chain_artifact(result, 'ground_truth')
                 # Emit integration test event once per test
                 try:
                     emit_integration_event_once(AnalyticsEvent.GT_ENCODER_INTEGRATION_TEST, {
@@ -2249,15 +3112,8 @@ def tensorleap_gt_encoder(name: str):
         inner.node_mapping = NodeMapping(name, NodeMappingType.GroundTruth)
 
         def mapping_inner(*args, **kwargs):
-            if _mapping_dataset_is_grouped:
-                class TempMapping:
-                    def __getitem__(self, key):
-                        # Indexing a grouped GT group (gt_list[i]) selects a sample from the same
-                        # ground-truth source; wiring is identical to the un-grouped GT node.
-                        return self
-            else:
-                class TempMapping:
-                    pass
+            class TempMapping:
+                pass
 
             ret = TempMapping()
             ret.node_mapping = mapping_inner.node_mapping
@@ -2534,12 +3390,12 @@ def tensorleap_status_table():
                     ready = False
                     next_step = row["name"]
 
-        if _crashed["value"]:
-            print(f"\nScript crashed before completing all steps. crashed at function '{_current_func['name']}'.")
+        if code_mapping_failure[0]:
+            print(f"\n{CROSS} {code_mapping_failure_mes}")
             return
 
-        if code_mapping_failure[0]:
-            print(f"\n{CROSS + code_mapping_failure_mes}.")
+        if _crashed["value"]:
+            print(f"\nScript crashed before completing all steps. crashed at function '{_current_func['name']}'.")
             return
 
         print(ready_mess) if ready else print(
@@ -2554,6 +3410,10 @@ def tensorleap_status_table():
             code_mapping_failure[0] = 1
         if status == "v":
             _set_status(name, CHECK)
+            # A decorator that reports success is no longer the "currently running"
+            # one, so clear the pointer. This keeps a later body-level raise from
+            # being blamed on the last decorator that actually succeeded.
+            _current_func["name"] = None
         else:
             _set_status(name, CROSS)
 
@@ -2565,6 +3425,12 @@ def tensorleap_status_table():
             table.append({"name": "tensorleap simulation (optional)", "Added to integration": UNKNOWN})
         _sim_tracking[sim_name] = "registered"
 
+    def register_decorator_row(name: str):
+        # Rows for decorators that only apply to some integration kinds (e.g. the autoregressive
+        # step hook) are appended on registration, so unrelated integrations never show them.
+        if _find_row(_remove_suffix(name, " (optional)")) is None:
+            table.append({"name": name, "Added to integration": UNKNOWN})
+
     def mark_sim_result(sim_name: str, passed: bool):
         _sim_tracking[sim_name] = "passed" if passed else "failed"
 
@@ -2600,10 +3466,14 @@ def tensorleap_status_table():
     def handle_exception(exc_type, exc_value, exc_traceback):
         _crashed["value"] = True
         crashed_name = _current_func["name"]
-        if crashed_name:
+        if crashed_name and not code_mapping_failure[0]:
             row = _find_row(crashed_name)
             if row:
                 row["Added to integration"] = CROSS
+        elif not crashed_name:
+            # Crash with no decorator currently running = the integration-test body /
+            # code flow itself failed; report that rather than blaming a decorator.
+            code_mapping_failure[0] = 1
 
         traceback.print_exception(exc_type, exc_value, exc_traceback)
         run_on_exit()
@@ -2611,11 +3481,12 @@ def tensorleap_status_table():
     atexit.register(run_on_exit)
     sys.excepthook = handle_exception
 
-    return set_current, update_env_params, register_sim, mark_sim_result
+    return set_current, update_env_params, register_sim, mark_sim_result, register_decorator_row
 
 
 if not _call_from_tl_platform:
-    set_current, update_env_params_func, register_sim, mark_sim_result = tensorleap_status_table()
+    set_current, update_env_params_func, register_sim, mark_sim_result, register_decorator_row = \
+        tensorleap_status_table()