code-loader 1.0.195.dev1__py3-none-any.whl → 1.0.197.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__)
 
@@ -26,13 +29,17 @@ from code_loader.contract.enums import MetricDirection, LeapDataType, DatasetMet
 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, LeapValidationError
+    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
+# Whether any integration test wrapper was entered in this process. The crash reporter uses it
+# to tell a raise in the test body (report as code-flow failure) from a module-level crash
+# before any test ran (report as a plain script crash).
+_integration_test_started = False
 _call_from_tl_platform = os.environ.get('IS_TENSORLEAP_PLATFORM') == 'true'
 
 # ---- warnings store (module-level) ----
@@ -206,50 +213,6 @@ def batch_warning(result, func_name):
         )
 
 
-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,16 +290,32 @@ 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],
-                                    func_name='integration_test', expected_names=["idx", "preprocess"], **kwargs)
+            # Keyword calls must work with the user function's OWN parameter names — the local
+            # run template calls integration_test(sample_id=..., preprocess=...) against a
+            # signature the user chose, not the canonical placeholder names.
+            expected_names = inspect.getfullargspec(integration_test_function)[0][:2]
+            if len(expected_names) < 2:
+                expected_names = ['idx', 'preprocess']
+            validate_args_structure(*args, types_order=[Union[int, str], PreprocessResponse],
+                                    func_name='integration_test', expected_names=expected_names,
+                                    **kwargs)
+            sample_id, preprocess_response = (
+                args[i] if i < len(args) else kwargs[name]
+                for i, name in enumerate(expected_names))
+            args, kwargs = (sample_id, preprocess_response), {}
             _validate_input_args(*args, **kwargs)
 
-            global _called_from_inside_tl_integration_test_decorator
+            global _called_from_inside_tl_integration_test_decorator, _integration_test_started
+            _integration_test_started = True
             # Clear integration test events for new test
             try:
                 clear_integration_events()
@@ -347,12 +326,12 @@ 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
@@ -370,7 +349,6 @@ def tensorleap_integration_test():
                                         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:
@@ -466,7 +444,7 @@ def tensorleap_load_model(prediction_types: Optional[List[PredictionTypeHandler]
             pass
 
         @lru_cache()
-        def inner(*args, **kwargs):
+        def _cached_inner(*args, **kwargs):
             if not _call_from_tl_platform:
                 set_current('tensorleap_load_model')
             validate_args_structure(*args, types_order=[],
@@ -624,14 +602,32 @@ def tensorleap_load_model(prediction_types: Optional[List[PredictionTypeHandler]
                 update_env_params_func("tensorleap_load_model", "v")
             return model_placeholder
 
+        def inner(*args, **kwargs):
+            model_placeholder = _cached_inner(*args, **kwargs)
+            # Re-assert on every call, outside the cache: each mapping rerun's load_model
+            # overwrites this global with its mapping placeholder, and on a cache hit the
+            # body above never runs — a second integration test would otherwise fail the
+            # model loop's identity check against the stale mapping placeholder.
+            global _last_loaded_model
+            _last_loaded_model = model_placeholder
+            return model_placeholder
+
         def mapping_inner():
             class ModelOutputPlaceholder:
                 def __init__(self):
                     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:
@@ -685,7 +681,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):
@@ -1562,10 +1561,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):
@@ -1583,20 +1585,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
@@ -1616,16 +1604,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
@@ -1635,52 +1621,38 @@ def tensorleap_metadata(
     return decorating_function
 
 
-def tensorleap_custom_latent_space():
+def tensorleap_custom_latent_space(name: Optional[str] = None):
     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')
+        ls_name = name if name is not None else user_function.__name__
+        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", ls_name, 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 '{ls_name}' 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
@@ -1692,7 +1664,7 @@ def tensorleap_custom_latent_space():
 
             return result
 
-        leap_binder.set_custom_latent_space(inner_without_validate)
+        leap_binder.set_custom_latent_space(inner_without_validate, ls_name)
 
         def inner(sample_id, preprocess_response):
             if os.environ.get(mapping_runtime_mode_env_var_mame):
@@ -1700,11 +1672,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
@@ -1712,18 +1682,263 @@ 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):
+    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 _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, preprocess) -> Optional[Dict[str, np.ndarray]]
-
-    First call per chain: prev_inputs=None, prev_outputs=None — returns the initial model inputs.
-    Later calls receive the previous step's model inputs and outputs (unbatched, per-sample) and
-    return the next step's full input dict; returning None ends the chain. Keys starting with '_'
-    are passthrough state: never fed to the model, handed back verbatim in prev_inputs next call.
-    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.
+        (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; any picklable nest of builtin/numpy values works (dicts, lists, tuples, sets,
+    numbers, strings, arrays — not user-defined classes, lambdas or open resources, since the
+    platform deserializes state without the integration code), 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
@@ -1742,62 +1957,95 @@ def tensorleap_autoregressive_step(latent_space_aggregation: str = 'last_step'):
         first_call_record: Dict[str, Any] = {'keys': None, 'shapes': None}
 
         def _normalize_args(args, kwargs):
-            expected_names = ["sample_id", "prev_inputs", "prev_outputs", "preprocess"]
-            validate_args_structure(
-                *args,
-                types_order=[Union[int, str], Union[dict, type(None)], Union[dict, type(None)],
-                             PreprocessResponse],
-                func_name=user_function.__name__, expected_names=expected_names, **kwargs)
+            expected_names = ["sample_id", "prev_inputs", "prev_outputs", "state", "preprocess"]
             normalized = []
-            for i, name in enumerate(expected_names):
-                normalized.append(args[i] if i < len(args) else kwargs[name])
+            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, preprocess_response):
+        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):
-            if result is None:
+            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) must return the initial '
-                     f'model inputs dict — returning None on the first call would produce an empty chain.')
+                     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(result, dict), \
+            assert isinstance(next_inputs, dict), \
                 (f'{user_function.__name__}() validation failed: '
-                 f'the return value must be a dict of named model input tensors, or None to end the '
-                 f'chain. Got {type(result)}.')
-            model_input_keys = [key for key in result if not key.startswith('_')]
-            assert model_input_keys, \
+                 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'the returned dict contains no model inputs — every key starts with "_" '
-                 f'(underscore keys are passthrough state, never fed to the model).')
-            for key, value in result.items():
+                 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 the returned dict must be strings. Got {type(key)}.')
-                if key.startswith('_'):
-                    continue
+                     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(result):
-            # Rules 2-3 of the autoregressive contract: identical key set and identical
-            # (non-passthrough) shapes on every call — compiled graphs need static shapes.
-            if result is None:
+        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(result.keys())
-            shapes = {key: tuple(value.shape) for key, value in result.items()
-                      if not key.startswith('_') and isinstance(value, np.ndarray)}
+            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
@@ -1819,36 +2067,26 @@ def tensorleap_autoregressive_step(latent_space_aggregation: str = 'last_step'):
                 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.')
-            if first_result is None or second_result is None:
-                assert first_result is None and second_result is None, error_message
-                return
-            assert set(first_result.keys()) == set(second_result.keys()), error_message
-            for key, value in first_result.items():
-                if key.startswith('_') or not isinstance(value, np.ndarray):
-                    continue
-                assert np.array_equal(value, second_result[key]), 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 not key.startswith('_') and 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, preprocess_response):
+            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 _nest_equal(value, second_inputs[key]), error_message
+            assert _nest_equal(first_state, second_state), error_message
+
+        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, preprocess_response)
+                result = user_function(sample_id, prev_inputs, prev_outputs, state,
+                                       preprocess_response)
             finally:
                 _called_from_inside_tl_decorator -= 1
             return result
@@ -1859,43 +2097,69 @@ def tensorleap_autoregressive_step(latent_space_aggregation: str = 'last_step'):
         def inner(*args, **kwargs):
             if not _call_from_tl_platform:
                 set_current('tensorleap_autoregressive_step')
-            sample_id, prev_inputs, prev_outputs, preprocess_response = _normalize_args(args, kwargs)
+            sample_id, prev_inputs, prev_outputs, state, preprocess_response = \
+                _normalize_args(args, kwargs)
 
-            is_top_level_in_test = (_called_from_inside_tl_decorator == 0
+            is_top_level = _called_from_inside_tl_decorator == 0
+            is_top_level_in_test = (is_top_level
                                     and _called_from_inside_tl_integration_test_decorator)
+            loop_context = _active_model_loop
+            if is_top_level_in_test and 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.')
+            # Outside the integration test a model loop may still be running (local debugging);
+            # the loop's phase machine must advance there too, or its model proxy always raises.
+            notify_loop = is_top_level and loop_context is not None
+            if notify_loop:
+                loop_context.on_hook_call(sample_id, prev_inputs, prev_outputs, state,
+                                          preprocess_response)
             if is_top_level_in_test:
-                prev_inputs = _squeeze_test_batch_dim(prev_inputs)
-                prev_outputs = _squeeze_test_batch_dim(prev_outputs)
+                # The wrapper adds the batch axis to returned inputs and model outputs come back
+                # batched; the hook body must see the same unbatched per-sample tensors it will
+                # see at engine runtime.
+                prev_inputs = _squeeze_leading_batch_axis(prev_inputs)
+                prev_outputs = _squeeze_leading_batch_axis(prev_outputs)
 
-            _validate_input_args(sample_id, prev_inputs, prev_outputs, preprocess_response)
+            _validate_input_args(sample_id, prev_inputs, prev_outputs, state, preprocess_response)
 
-            result = inner_without_validate(sample_id, prev_inputs, prev_outputs, 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,
-                                                       preprocess_response)
+                                                       state_for_second_call, preprocess_response)
                 _assert_deterministic(result, second_result)
 
             _validate_result(result, is_first_call=prev_inputs is None)
-            _validate_consistency(result)
+            next_inputs, new_state = result
+            _validate_consistency(next_inputs)
+            _validate_state_types(new_state)
 
-            if is_top_level_in_test and result is not None:
-                result = {key: (np.expand_dims(value, axis=0)
-                                if not key.startswith('_') and isinstance(value, np.ndarray) else value)
-                          for key, value in result.items()}
+            if is_top_level_in_test and next_inputs is not None:
+                next_inputs = {key: np.expand_dims(value, axis=0)
+                               for key, value in next_inputs.items()}
+            if notify_loop:
+                loop_context.on_hook_return(next_inputs, new_state)
 
             if not _call_from_tl_platform:
                 update_env_params_func("tensorleap_autoregressive_step", "v")
-            return result
+            return next_inputs, new_state
 
         class _MappingInputsPlaceholder:
-            """Mapping-mode stand-in for the hook's returned dict.
+            """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 non-passthrough key, once.
+            NodeConnection is recorded per key, once.
             """
             def __init__(self):
                 self._items: Dict[str, Any] = {}
@@ -1908,8 +2172,7 @@ def tensorleap_autoregressive_step(latent_space_aggregation: str = 'last_step'):
                         pass
                     ret = TempMapping()
                     ret.node_mapping = NodeMapping(key, NodeMappingType.Input)
-                    if not key.startswith('_'):
-                        leap_binder.mapping_connections.append(NodeConnection(ret.node_mapping, None))
+                    leap_binder.mapping_connections.append(NodeConnection(ret.node_mapping, None))
                     self._items[key] = ret
                 return self._items[key]
 
@@ -1934,9 +2197,17 @@ def tensorleap_autoregressive_step(latent_space_aggregation: str = 'last_step'):
                 self._unsupported_iteration()
 
         mapping_placeholder = _MappingInputsPlaceholder()
+        mapping_state_placeholder = _MappingStatePlaceholder()
 
         def mapping_inner(*args, **kwargs):
-            return mapping_placeholder
+            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):
@@ -1952,6 +2223,396 @@ def tensorleap_autoregressive_step(latent_space_aggregation: str = 'last_step'):
     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
+            in_test = bool(_called_from_inside_tl_integration_test_decorator)
+            if _active_model_loop is not None:
+                _fail('model loops cannot nest or run concurrently.')
+            if in_test and _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
+            if in_test:
+                _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)
@@ -1982,15 +2643,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:
@@ -2328,10 +2980,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)}.')
@@ -2341,13 +2997,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
@@ -2364,27 +3013,18 @@ 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 _called_from_inside_tl_decorator == 0 and _called_from_inside_tl_integration_test_decorator:
-                if grouped:
-                    # A grouped result is a list of per-sample arrays (never stacked, so a
-                    # B=1-only model can be fed one sample at a time). Add a batch dim to each
-                    # sample, mirroring the flat path's single-sample expand_dims.
-                    if isinstance(result, list):
-                        result = [np.expand_dims(r, axis=0) for r in result]
-                else:
-                    batch_warning(result, user_function.__name__)
-                    result = np.expand_dims(result, axis=0)
+                batch_warning(result, user_function.__name__)
+                result = np.expand_dims(result, axis=0)
                 # Emit integration test event once per test
                 try:
                     emit_integration_event_once(AnalyticsEvent.INPUT_ENCODER_INTEGRATION_TEST, {
@@ -2405,15 +3045,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
@@ -2443,10 +3076,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)}.')
@@ -2454,14 +3092,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
@@ -2478,27 +3108,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 _called_from_inside_tl_decorator == 0 and _called_from_inside_tl_integration_test_decorator:
-                if grouped:
-                    # A grouped result is a list of per-sample arrays (never stacked, so a
-                    # B=1-only model can be fed one sample at a time). Add a batch dim to each
-                    # sample, mirroring the flat path's single-sample expand_dims.
-                    if isinstance(result, list):
-                        result = [np.expand_dims(r, axis=0) for r in result]
-                else:
-                    batch_warning(result, user_function.__name__)
-                    result = np.expand_dims(result, axis=0)
+                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, {
@@ -2513,15 +3135,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
@@ -2772,6 +3387,12 @@ def tensorleap_status_table():
         if not started_from("leap_integration.py"):
             return
 
+        if leap_binder.setup_container.autoregressive_step is not None:
+            # Input encoders are forbidden in autoregressive integrations — the step hook is
+            # the sole source of model inputs — so the row is not applicable: it must not be
+            # printed, gate readiness, or be recommended as a next step.
+            table[:] = [row for row in table if row["name"] != "tensorleap_input_encoder"]
+
         ready_mess = "\nAll parts have been successfully set. If no errors accured, you can now push the project to the Tensorleap system."
         not_ready_mess = "\nSome mandatory components have not yet been added to the Integration test. Recommended next interface to add is: "
         mandatory_ready_mess = "\nAll mandatory parts have been successfully set. If no errors accured, you can now push the project to the Tensorleap system or continue to the next optional reccomeded interface,adding: "
@@ -2803,7 +3424,11 @@ def tensorleap_status_table():
             return
 
         if _crashed["value"]:
-            print(f"\nScript crashed before completing all steps. crashed at function '{_current_func['name']}'.")
+            crashed_at = _current_func['name']
+            if crashed_at:
+                print(f"\nScript crashed before completing all steps. crashed at function '{crashed_at}'.")
+            else:
+                print("\nScript crashed before completing all steps.")
             return
 
         print(ready_mess) if ready else print(
@@ -2878,9 +3503,11 @@ def tensorleap_status_table():
             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.
+        elif not crashed_name and _integration_test_started:
+            # Crash with no decorator currently running while an integration test has been
+            # entered = the test body / code flow itself failed; report that rather than
+            # blaming a decorator. Without a started test (a module-level crash between
+            # decorator registrations) the plain script-crash message is the truthful one.
             code_mapping_failure[0] = 1
 
         traceback.print_exception(exc_type, exc_value, exc_traceback)