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

code_loader/inner_leap_binder/leapbinder_decorators.py

@@ -1,6 +1,4 @@
 # mypy: ignore-errors
-import copy
-import hashlib
 import inspect
 import os
 import warnings
@@ -15,8 +13,7 @@ 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, \
-    validate_autoregressive_state_types, autoregressive_nests_equal
+from code_loader.utils import get_metadata_type_from_variable, map_dict_to_metadata_types
 
 logger = logging.getLogger(__name__)
 
@@ -29,17 +26,13 @@ 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, \
-    map_leap_data_type_to_visualizer_class
+    LeapHorizontalBar, LeapImageWithBBox, LeapImageWithHeatmap, LeapVideo, LeapValidationError
 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
-# 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
+_mapping_dataset_is_grouped = False
 _call_from_tl_platform = os.environ.get('IS_TENSORLEAP_PLATFORM') == 'true'
 
 # ---- warnings store (module-level) ----
@@ -213,6 +206,50 @@ 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)]
@@ -290,21 +327,16 @@ def tensorleap_integration_test():
 
         def _validate_input_args(*args, **kwargs):
             sample_id, preprocess_response = args
-            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."
-            )
+            _validate_id_or_group(sample_id, preprocess_response, 'tensorleap_integration_test')
 
         def inner(*args, **kwargs):
             if not _call_from_tl_platform:
                 set_current('tensorleap_integration_test')
-            validate_args_structure(*args, types_order=[Union[int, str], PreprocessResponse],
+            validate_args_structure(*args, types_order=[Union[int, str, list], PreprocessResponse],
                                     func_name='integration_test', expected_names=["idx", "preprocess"], **kwargs)
             _validate_input_args(*args, **kwargs)
 
-            global _called_from_inside_tl_integration_test_decorator, _integration_test_started
-            _integration_test_started = True
+            global _called_from_inside_tl_integration_test_decorator
             # Clear integration test events for new test
             try:
                 clear_integration_events()
@@ -315,12 +347,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
@@ -338,6 +370,7 @@ 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:
@@ -587,8 +620,6 @@ 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
@@ -599,16 +630,8 @@ def tensorleap_load_model(prediction_types: Optional[List[PredictionTypeHandler]
                     self.node_mapping = NodeMapping('', NodeMappingType.Prediction0)
 
                 def __getitem__(self, key):
-                    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)
+                    assert isinstance(key, int), \
+                        f'Expected key to be an int, got {type(key)} instead.'
 
                     ret = TempMapping()
                     try:
@@ -662,10 +685,7 @@ def tensorleap_load_model(prediction_types: Optional[List[PredictionTypeHandler]
 
                     return FollowInputIndex()
 
-            model_placeholder = ModelPlaceholder()
-            global _last_loaded_model
-            _last_loaded_model = model_placeholder
-            return model_placeholder
+            return ModelPlaceholder()
 
         def final_inner(*args, **kwargs):
             if os.environ.get(mapping_runtime_mode_env_var_mame):
@@ -1542,13 +1562,10 @@ 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], 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_input_args(sample_id: Union[int, str, list], preprocess_response: PreprocessResponse):
+            _validate_id_or_group(sample_id, preprocess_response, user_function.__name__)
 
-        def _validate_result(result):
+        def _validate_single(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):
@@ -1566,6 +1583,20 @@ 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
@@ -1585,14 +1616,16 @@ 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], PreprocessResponse],
+            validate_args_structure(*args, types_order=[Union[int, str, list], 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)
+            _validate_result(result, grouped, group_size)
             if not _call_from_tl_platform:
                 update_env_params_func("tensorleap_metadata", "v")
             return result
@@ -1604,35 +1637,50 @@ def tensorleap_metadata(
 
 def tensorleap_custom_latent_space():
     def decorating_function(user_function: SectionCallableInterface):
-        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_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_result(result):
-            assert isinstance(result, np.ndarray), \
+        def _validate_single(single_result):
+            assert isinstance(single_result, np.ndarray), \
                 (f'tensorleap_custom_latent_space validation failed: '
-                 f'The return type should be a numpy array. Got {type(result)}.')
-            if result.ndim > 1:
-                flat_dim = int(np.prod(result.shape))
+                 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))
                 store_general_warning(
-                    key=("tensorleap_custom_latent_space_flatten", tuple(result.shape)),
+                    key=("tensorleap_custom_latent_space_flatten", tuple(single_result.shape)),
                     message=(
-                        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"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"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
@@ -1652,9 +1700,11 @@ 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)
+            _validate_result(result, grouped, group_size)
             return result
 
         return inner
@@ -1662,260 +1712,18 @@ 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.
+        (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.
 
     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
@@ -1934,95 +1742,62 @@ 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", "state", "preprocess"]
+            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)
             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__}.')
+            for i, name in enumerate(expected_names):
+                normalized.append(args[i] if i < len(args) else kwargs[name])
             return normalized
 
-        def _validate_input_args(sample_id, prev_inputs, prev_outputs, state, preprocess_response):
+        def _validate_input_args(sample_id, prev_inputs, prev_outputs, 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:
+            if result 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.')
+                     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.')
                 return
-            assert isinstance(next_inputs, dict), \
+            assert isinstance(result, 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'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'{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():
+                 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():
                 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.')
+                     f'all keys in the returned dict must be strings. Got {type(key)}.')
+                if key.startswith('_'):
+                    continue
                 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:
+        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:
                 return
-            keys = set(next_inputs.keys())
-            shapes = {key: tuple(value.shape) for key, value in next_inputs.items()}
+            keys = set(result.keys())
+            shapes = {key: tuple(value.shape) for key, value in result.items()
+                      if not key.startswith('_') and isinstance(value, np.ndarray)}
             if first_call_record['keys'] is None:
                 first_call_record['keys'] = keys
                 first_call_record['shapes'] = shapes
@@ -2044,19 +1819,14 @@ 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.')
-            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
+            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
@@ -2067,18 +1837,18 @@ def tensorleap_autoregressive_step(latent_space_aggregation: str = 'last_step'):
                 return None
             squeezed = {}
             for key, value in tensors_dict.items():
-                if isinstance(value, np.ndarray) and value.ndim > 0 and value.shape[0] == 1:
+                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, state, preprocess_response):
+        def inner_without_validate(sample_id, prev_inputs, prev_outputs, 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)
+                result = user_function(sample_id, prev_inputs, prev_outputs, preprocess_response)
             finally:
                 _called_from_inside_tl_decorator -= 1
             return result
@@ -2089,66 +1859,43 @@ 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, state, preprocess_response = \
-                _normalize_args(args, kwargs)
+            sample_id, prev_inputs, prev_outputs, preprocess_response = _normalize_args(args, kwargs)
 
-            is_top_level = _called_from_inside_tl_decorator == 0
-            is_top_level_in_test = (is_top_level
+            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 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)
 
-            _validate_input_args(sample_id, prev_inputs, prev_outputs, state, preprocess_response)
+            _validate_input_args(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)
+            result = inner_without_validate(sample_id, prev_inputs, prev_outputs, 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)
+                                                       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)
+            _validate_consistency(result)
 
-            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 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 not _call_from_tl_platform:
                 update_env_params_func("tensorleap_autoregressive_step", "v")
-            return next_inputs, new_state
+            return result
 
         class _MappingInputsPlaceholder:
-            """Mapping-mode stand-in for the hook's returned inputs dict.
+            """Mapping-mode stand-in for the hook's returned 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.
+            NodeConnection is recorded per non-passthrough key, once.
             """
             def __init__(self):
                 self._items: Dict[str, Any] = {}
@@ -2161,7 +1908,8 @@ def tensorleap_autoregressive_step(latent_space_aggregation: str = 'last_step'):
                         pass
                     ret = TempMapping()
                     ret.node_mapping = NodeMapping(key, NodeMappingType.Input)
-                    leap_binder.mapping_connections.append(NodeConnection(ret.node_mapping, None))
+                    if not key.startswith('_'):
+                        leap_binder.mapping_connections.append(NodeConnection(ret.node_mapping, None))
                     self._items[key] = ret
                 return self._items[key]
 
@@ -2186,17 +1934,9 @@ 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):
-            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)
+            return mapping_placeholder
 
         def final_inner(*args, **kwargs):
             if os.environ.get(mapping_runtime_mode_env_var_mame):
@@ -2212,396 +1952,6 @@ 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)
@@ -2632,6 +1982,15 @@ 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:
@@ -2969,14 +2328,10 @@ 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], 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_input_args(sample_id: Union[int, str, list], preprocess_response: PreprocessResponse):
+            _validate_id_or_group(sample_id, preprocess_response, user_function.__name__)
 
-        def _validate_result(result):
-            validate_output_structure(result, func_name=user_function.__name__, expected_type_name="np.ndarray")
+        def _validate_single(result):
             assert isinstance(result, np.ndarray), \
                 (f'{user_function.__name__}() validation failed: '
                  f'Unsupported return type. Should be a numpy array. Got {type(result)}.')
@@ -2986,6 +2341,13 @@ 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
@@ -3002,18 +2364,27 @@ 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], PreprocessResponse],
+            validate_args_structure(*args, types_order=[Union[int, str, list], 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)
+            _validate_result(result, grouped, group_size)
 
             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)
+                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)
                 # Emit integration test event once per test
                 try:
                     emit_integration_event_once(AnalyticsEvent.INPUT_ENCODER_INTEGRATION_TEST, {
@@ -3034,8 +2405,15 @@ 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):
-            class TempMapping:
-                pass
+            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
 
             ret = TempMapping()
             ret.node_mapping = mapping_inner.node_mapping
@@ -3065,15 +2443,10 @@ 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], 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_input_args(sample_id: Union[int, str, list], preprocess_response: PreprocessResponse):
+            _validate_id_or_group(sample_id, preprocess_response, user_function.__name__)
 
-        def _validate_result(result):
-            validate_output_structure(result, func_name=user_function.__name__, expected_type_name="np.ndarray",
-                                      gt_flag=True)
+        def _validate_single(result):
             assert isinstance(result, np.ndarray), \
                 (f'{user_function.__name__}() validation failed: '
                  f'Unsupported return type. Should be a numpy array. Got {type(result)}.')
@@ -3081,6 +2454,14 @@ 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
@@ -3097,19 +2478,27 @@ 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], PreprocessResponse],
+            validate_args_structure(*args, types_order=[Union[int, str, list], 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)
+            _validate_result(result, grouped, group_size)
 
             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')
+                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)
                 # Emit integration test event once per test
                 try:
                     emit_integration_event_once(AnalyticsEvent.GT_ENCODER_INTEGRATION_TEST, {
@@ -3124,8 +2513,15 @@ def tensorleap_gt_encoder(name: str):
         inner.node_mapping = NodeMapping(name, NodeMappingType.GroundTruth)
 
         def mapping_inner(*args, **kwargs):
-            class TempMapping:
-                pass
+            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
 
             ret = TempMapping()
             ret.node_mapping = mapping_inner.node_mapping
@@ -3407,11 +2803,7 @@ def tensorleap_status_table():
             return
 
         if _crashed["value"]:
-            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.")
+            print(f"\nScript crashed before completing all steps. crashed at function '{_current_func['name']}'.")
             return
 
         print(ready_mess) if ready else print(
@@ -3486,11 +2878,9 @@ def tensorleap_status_table():
             row = _find_row(crashed_name)
             if row:
                 row["Added to integration"] = CROSS
-        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.
+        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)