code-loader 1.0.188.dev1__py3-none-any.whl → 1.0.189.dev0__py3-none-any.whl

code_loader/contract/datasetclasses.py

@@ -92,6 +92,16 @@ SectionCallableInterface = Callable[[Union[int, str], PreprocessResponse], npt.N
 InstanceCallableInterface = Callable[[Union[int, str], PreprocessResponse, int], Optional[ElementInstance]]
 InstanceLengthCallableInterface = Callable[[Union[int, str], PreprocessResponse], int]
 
+# (sample_id, prev_inputs, prev_outputs, preprocess) -> next model inputs, or None to end the chain.
+# First call per chain receives prev_inputs=None, prev_outputs=None and returns the initial inputs.
+# Keys starting with '_' are passthrough state: never fed to the model, handed back verbatim in
+# prev_inputs on the next call. The hook must be stateless and deterministically seeded (the engine
+# may replay a step after crash recovery and relies on identical results).
+AutoregressiveStepCallableInterface = Callable[
+    [Union[int, str], Optional[Dict[str, npt.NDArray[np.float32]]], Optional[Dict[str, npt.NDArray[np.float32]]],
+     PreprocessResponse],
+    Optional[Dict[str, npt.NDArray[np.float32]]]]
+
 MetadataSectionCallableInterface = Union[
     Callable[[Union[int, str], PreprocessResponse], int],
     Callable[[Union[int, str], PreprocessResponse], Dict[str, int]],
@@ -239,6 +249,15 @@ class CustomLatentSpaceHandler:
     function: SectionCallableInterface
     name: str = 'custom_latent_space'
 
+
+@dataclass
+class AutoregressiveStepHandler:
+    function: AutoregressiveStepCallableInterface
+    name: str = 'autoregressive_step'
+    # Per-model-input shapes (unbatched, underscore passthrough keys excluded), discovered by running
+    # the hook's first call at parse time. Fills the role InputHandler.shape plays for input encoders.
+    input_shapes: Optional[Dict[str, List[int]]] = None
+
 @dataclass
 class PredictionTypeHandler:
     name: str
@@ -279,6 +298,7 @@ class DatasetIntegrationSetup:
     custom_layers: Dict[str, CustomLayerHandler] = field(default_factory=dict)
     custom_latent_space: Optional[CustomLatentSpaceHandler] = None
     simulations: List[SimulationHandler] = field(default_factory=list)
+    autoregressive_step: Optional[AutoregressiveStepHandler] = None
 
 
 @dataclass

code_loader/inner_leap_binder/leapbinder.py

@@ -1,6 +1,6 @@
 
 import inspect
-from typing import Callable, List, Optional, Dict, Any, Type, Union
+from typing import Callable, List, Optional, Dict, Any, Type, Union, get_args
 
 import numpy as np
 import numpy.typing as npt
@@ -13,7 +13,7 @@ from code_loader.contract.datasetclasses import SectionCallableInterface, InputH
     CustomMultipleReturnCallableInterfaceMultiArgs, DatasetBaseHandler, custom_latent_space_attribute, \
     RawInputsForHeatmap, VisualizerHandlerData, MetricHandlerData, CustomLossHandlerData, SamplePreprocessResponse, \
     ElementInstanceMasksHandler, InstanceCallableInterface, CustomLatentSpaceHandler, InstanceMetricHandler, \
-    SimulationHandler, _simulation_context
+    SimulationHandler, _simulation_context, AutoregressiveStepHandler, AutoregressiveStepCallableInterface
 from code_loader.contract.enums import LeapDataType, DataStateEnum, DataStateType, MetricDirection, DatasetMetadataType
 from code_loader.contract.mapping import NodeConnection, NodeMapping, NodeMappingType
 from code_loader.contract.responsedataclasses import DatasetTestResultPayload, LeapAnalysisConfiguration
@@ -31,6 +31,27 @@ from code_loader.visualizers.default_visualizers import DefaultVisualizer, \
 mapping_runtime_mode_env_var_mame = '__MAPPING_RUNTIME_MODE__'
 
 
+def _stringized_annotation_type_name(annotation: Any) -> Optional[str]:
+    """Bare type name if ``annotation`` is a stringized annotation (from ``from __future__
+    import annotations`` or a quoted hint), else ``None``. code_loader inspects raw
+    annotations expecting real classes, so stringization silently breaks type detection."""
+    if isinstance(annotation, str):
+        return annotation.split("[")[0].strip().rsplit(".", 1)[-1]
+    return None
+
+
+def _reject_stringized_sample_preprocess_response(function: Callable[..., Any], arg_name: str, annotation: Any) -> None:
+    """Fail fast on a stringized SamplePreprocessResponse arg (otherwise it silently
+    mis-wires as a regular input, surfacing only on the platform)."""
+    if _stringized_annotation_type_name(annotation) == SamplePreprocessResponse.__name__:
+        raise Exception(
+            f"Argument '{arg_name}' of function '{function.__name__}' is annotated with a string "
+            f"('{annotation}') instead of the SamplePreprocessResponse type. This usually means the "
+            f"file uses 'from __future__ import annotations' (or a quoted hint), which stringizes "
+            f"annotations and breaks Tensorleap type detection. Remove that import (or the quotes) so "
+            f"SamplePreprocessResponse is referenced as a real type.")
+
+
 
 
 class LeapBinder:
@@ -123,6 +144,7 @@ class LeapBinder:
         regular_arg_names = inspect.getfullargspec(function)[0]
         preprocess_response_arg_name = None
         for arg_name, arg_type in inspect.getfullargspec(function).annotations.items():
+            _reject_stringized_sample_preprocess_response(function, arg_name, arg_type)
             if arg_type == SamplePreprocessResponse:
                 if preprocess_response_arg_name is not None:
                     raise Exception("only one argument can be of type SamplePreprocessResponse")
@@ -144,8 +166,8 @@ class LeapBinder:
 
         if visualizer_type.value not in map_leap_data_type_to_visualizer_class:
             raise Exception(
-                f'The visualizer_type is invalid. current visualizer_type: {visualizer_type}, '  # type: ignore[attr-defined]
-                f'should be one of : {", ".join([arg.__name__ for arg in LeapData.__args__])}')
+                f'The visualizer_type is invalid. current visualizer_type: {visualizer_type}, '
+                f'should be one of : {", ".join([arg.__name__ for arg in get_args(LeapData)])}')
 
         func_annotations = function.__annotations__
         if "return" not in func_annotations:
@@ -155,10 +177,17 @@ class LeapBinder:
                   f"https://docs.python.org/3/library/typing.html")
         else:
             return_type = func_annotations["return"]
-            if return_type not in LeapData.__args__:  # type: ignore[attr-defined]
+            if isinstance(return_type, str):
+                raise Exception(
+                    f"The return type hint of function '{function.__name__}' is a string "
+                    f"('{return_type}') instead of a type. This usually means the file uses "
+                    f"'from __future__ import annotations' (or a quoted return hint), which stringizes "
+                    f"annotations and prevents Tensorleap from detecting the visualizer type. Remove "
+                    f"that import (or the quotes) so the return type is a real class.")
+            if return_type not in get_args(LeapData):
                 raise Exception(
-                    f'The return type of function {function.__name__} is invalid. current return type: {return_type}, '  # type: ignore[attr-defined]
-                    f'should be one of : {", ".join([arg.__name__ for arg in LeapData.__args__])}')
+                    f'The return type of function {function.__name__} is invalid. current return type: {return_type}, '
+                    f'should be one of : {", ".join([arg.__name__ for arg in get_args(LeapData)])}')
 
             expected_return_type = map_leap_data_type_to_visualizer_class[visualizer_type.value]
             if not issubclass(return_type, expected_return_type):
@@ -296,6 +325,7 @@ class LeapBinder:
         regular_arg_names = inspect.getfullargspec(function)[0]
         preprocess_response_arg_name = None
         for arg_name, arg_type in inspect.getfullargspec(function).annotations.items():
+            _reject_stringized_sample_preprocess_response(function, arg_name, arg_type)
             if arg_type == SamplePreprocessResponse:
                 if preprocess_response_arg_name is not None:
                     raise Exception("only one argument can be of type SamplePreprocessResponse")
@@ -339,6 +369,7 @@ class LeapBinder:
         regular_arg_names = inspect.getfullargspec(function)[0]
         preprocess_response_arg_name = None
         for arg_name, arg_type in inspect.getfullargspec(function).annotations.items():
+            _reject_stringized_sample_preprocess_response(function, arg_name, arg_type)
             if arg_type == SamplePreprocessResponse:
                 if preprocess_response_arg_name is not None:
                     raise Exception("only one argument can be of type SamplePreprocessResponse")
@@ -380,6 +411,7 @@ class LeapBinder:
         regular_arg_names = inspect.getfullargspec(function)[0]
         preprocess_response_arg_name = None
         for arg_name, arg_type in inspect.getfullargspec(function).annotations.items():
+            _reject_stringized_sample_preprocess_response(function, arg_name, arg_type)
             if arg_type == SamplePreprocessResponse:
                 if preprocess_response_arg_name is not None:
                     raise Exception("only one argument can be of type SamplePreprocessResponse")
@@ -492,6 +524,26 @@ class LeapBinder:
         """
         self.setup_container.custom_latent_space = CustomLatentSpaceHandler(function)
 
+    def set_autoregressive_step(self, function: AutoregressiveStepCallableInterface) -> None:
+        """
+        Set the autoregressive step hook — the feedback function that drives a chain:
+        it supplies the model's initial inputs on its first call (prev_inputs=None, prev_outputs=None)
+        and turns each step's inputs/outputs into the next step's inputs until it returns None.
+        An autoregressive integration has no input encoders; the hook is the sole input source.
+        """
+        if self.setup_container.autoregressive_step is not None:
+            raise Exception('tensorleap_autoregressive_step is already defined. '
+                            'Only one autoregressive step hook is allowed per integration.')
+        self.setup_container.autoregressive_step = AutoregressiveStepHandler(function)
+
+        # Builtin chain metadata, declared at parse time so it survives the reporter's
+        # metadata type mapping; the placeholder values are overwritten by the engine when a
+        # chain finalizes (realized length, truncated-by-safety-cap flag).
+        def builtin_chain_metadata(idx: Any, preprocess: PreprocessResponse) -> Dict[str, Any]:
+            return {'length': -1, 'truncated': False}
+
+        self.set_metadata(builtin_chain_metadata, 'builtin_chain')
+
     def set_custom_layer(self, custom_layer: Type[Any], name: str, inspect_layer: bool = False,
                          kernel_index: Optional[int] = None, use_custom_latent_space: bool = False) -> None:
         """
@@ -702,6 +754,7 @@ class LeapBinder:
                     )
 
     def check(self) -> None:
+        self.validate_autoregressive_setup()
         preprocess_result = self.get_preprocess_result()
         self.check_preprocess(preprocess_result)
         self.check_handlers(preprocess_result)
@@ -709,6 +762,32 @@ class LeapBinder:
         self.validate_ignore_latent_spaces()
         print("Successful!")
 
+    def validate_autoregressive_setup(self) -> None:
+        """Structural rules for autoregressive integrations (no-op when the hook is absent).
+
+        Order-independent: runs after the whole script registered, because decorators may
+        appear in any order in the integration file.
+        """
+        if self.setup_container.autoregressive_step is None:
+            return
+        if self.setup_container.inputs:
+            raise Exception(
+                "An autoregressive integration must not declare input encoders — "
+                "tensorleap_autoregressive_step is the sole source of model inputs "
+                "(its first call returns the initial inputs). Found input encoders: "
+                f"{[handler.name for handler in self.setup_container.inputs]}.")
+        if self.setup_container.instance_masks:
+            raise Exception(
+                "Element instances are not supported together with tensorleap_autoregressive_step: "
+                "instance generation masks the sample's encoded inputs, which do not exist in an "
+                "autoregressive integration. Remove the instance encoders or the autoregressive hook.")
+        if self.setup_container.custom_latent_space is not None:
+            raise Exception(
+                "tensorleap_custom_latent_space is not supported together with "
+                "tensorleap_autoregressive_step: the custom latent space function only sees "
+                "(sample_id, preprocess) — it runs before generation and cannot observe the chain. "
+                "The chain's latent is the final step's model latent.")
+
     def validate_ignore_latent_spaces(self) -> None:
         """Validate leap_analysis_configuration.ignore_latent_spaces against this binder.
 

code_loader/inner_leap_binder/leapbinder_decorators.py

@@ -21,7 +21,7 @@ from code_loader.contract.datasetclasses import CustomCallableInterfaceMultiArgs
     CustomMultipleReturnCallableInterfaceMultiArgs, ConfusionMatrixCallableInterfaceMultiArgs, CustomCallableInterface, \
     VisualizerCallableInterface, MetadataSectionCallableInterface, PreprocessResponse, SectionCallableInterface, \
     ConfusionMatrixElement, SamplePreprocessResponse, PredictionTypeHandler, InstanceCallableInterface, ElementInstance, \
-    InstanceLengthCallableInterface
+    InstanceLengthCallableInterface, AutoregressiveStepCallableInterface
 from code_loader.contract.enums import MetricDirection, LeapDataType, DatasetMetadataType, DataStateType
 from code_loader import leap_binder, LeapLoader
 from code_loader.contract.mapping import NodeMapping, NodeMappingType, NodeConnection
@@ -223,6 +223,59 @@ def _add_mapping_connections(connects_to, arg_names, node_mapping_type, name):
         _add_mapping_connection(user_unique_name, connection_destinations, arg_names, name, node_mapping_type)
 
 
+def _require_sample_preprocess_response_supplied(user_function: Callable, args: tuple, kwargs: dict) -> None:
+    """A SamplePreprocessResponse argument is auto-injected by the platform / check_dataset
+    but NOT inside integration_test, where the author calls the function directly. Fail fast
+    with an actionable message instead of a raw 'missing argument' TypeError."""
+    spr_arg_name = None
+    for arg_name, arg_type in inspect.getfullargspec(user_function).annotations.items():
+        if arg_type == SamplePreprocessResponse:
+            spr_arg_name = arg_name
+            break
+    if spr_arg_name is None:
+        return
+    signature = inspect.signature(user_function)
+    param = signature.parameters.get(spr_arg_name)
+    # Only enforce a REQUIRED SamplePreprocessResponse arg. If it has a default it is
+    # optional and the function may legitimately be called with or without it.
+    if param is None or param.default is not inspect.Parameter.empty:
+        return
+    try:
+        supplied = signature.bind_partial(*args, **kwargs).arguments
+    except TypeError:
+        return
+    if spr_arg_name not in supplied:
+        raise Exception(
+            f"'{user_function.__name__}' takes a SamplePreprocessResponse argument "
+            f"'{spr_arg_name}', but it was not provided. Inside integration_test pass it "
+            f"explicitly, e.g. {user_function.__name__}(..., "
+            f"SamplePreprocessResponse(sample_id, preprocess)).")
+
+
+def _warn_loss_inputs_not_wired_to_model() -> None:
+    """Warn when a custom loss consumes an input encoder the model is not given.
+
+    Such an input stays a plain NodeMappingType.Input (model-fed inputs get re-tagged to
+    Input0/Input1/...) and crashes the platform model build. Losses only — the same pattern
+    in a metric or visualizer builds fine."""
+    loss_types = {NodeMappingType.Loss, NodeMappingType.CustomLoss}
+    for connection in leap_binder.mapping_connections:
+        if connection.node is None or connection.node.type not in loss_types:
+            continue
+        for arg_name, source in (connection.node_inputs or {}).items():
+            if source is not None and source.type == NodeMappingType.Input:
+                input_name = source.name or arg_name
+                store_general_warning(
+                    key=("loss_input_not_wired_to_model", connection.node.name, input_name),
+                    message=(
+                        f"Custom loss '{connection.node.name}' consumes input '{input_name}', "
+                        f"which is not fed to the model. This is not supported: the platform builds "
+                        f"losses into the model graph, so an input the model does not take has no "
+                        f"source there and the model build fails. Instead, carry this data in the "
+                        f"ground truth and derive it inside the loss (not as a separate input "
+                        f"encoder)."))
+
+
 def tensorleap_integration_test():
     def decorating_function(integration_test_function: Callable):
         leap_binder.integration_test_func = integration_test_function
@@ -279,6 +332,8 @@ def tensorleap_integration_test():
             finally:
                 _called_from_inside_tl_integration_test_decorator = False
 
+            _warn_loss_inputs_not_wired_to_model()
+
             leap_binder.check()
 
         return inner
@@ -822,6 +877,7 @@ def tensorleap_custom_metric(name: str,
             _called_from_inside_tl_decorator += 1
 
             try:
+                _require_sample_preprocess_response_supplied(user_function, args, kwargs)
                 result = user_function(*args, **kwargs)
             finally:
                 _called_from_inside_tl_decorator -= 1
@@ -1031,6 +1087,7 @@ def tensorleap_custom_instances_metric(name: str,
             _called_from_inside_tl_decorator += 1
 
             try:
+                _require_sample_preprocess_response_supplied(user_function, args, kwargs)
                 result = user_function(*args, **kwargs)
             finally:
                 _called_from_inside_tl_decorator -= 1
@@ -1383,6 +1440,7 @@ def tensorleap_custom_visualizer(name: str, visualizer_type: LeapDataType,
             _called_from_inside_tl_decorator += 1
 
             try:
+                _require_sample_preprocess_response_supplied(user_function, args, kwargs)
                 result = user_function(*args, **kwargs)
             finally:
                 _called_from_inside_tl_decorator -= 1
@@ -1570,6 +1628,232 @@ def tensorleap_custom_latent_space():
     return decorating_function
 
 
+def tensorleap_autoregressive_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.
+    """
+    def decorating_function(user_function: AutoregressiveStepCallableInterface):
+        # Contract: shapes and key set are fixed across all calls AND all chains (compiled graphs
+        # need static shapes), so one record shared across samples is the correct scope.
+        first_call_record: Dict[str, Any] = {'keys': None, 'shapes': None}
+
+        def _normalize_args(args, kwargs):
+            expected_names = ["sample_id", "prev_inputs", "prev_outputs", "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, 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, 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 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 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.')
+                return
+            assert isinstance(result, 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'{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():
+                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
+                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:
+                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)}
+            if first_call_record['keys'] is None:
+                first_call_record['keys'] = keys
+                first_call_record['shapes'] = shapes
+                return
+            assert keys == first_call_record['keys'], \
+                (f'{user_function.__name__}() validation failed: '
+                 f'every call must return the same key set as the first call. '
+                 f'First call keys: {sorted(first_call_record["keys"])}, this call: {sorted(keys)}.')
+            for key, shape in shapes.items():
+                assert shape == first_call_record['shapes'][key], \
+                    (f'{user_function.__name__}() validation failed: '
+                     f'tensor shapes must be identical on every call (fixed-length padding + mask for '
+                     f'growing sequences). Input "{key}" was {first_call_record["shapes"][key]} on the '
+                     f'first call and {shape} now.')
+
+        def _assert_deterministic(first_result, second_result):
+            error_message = (
+                f'{user_function.__name__}() validation failed: '
+                f'the hook is not deterministic — two calls with identical arguments returned different '
+                f'outputs. Seed any stochasticity (e.g. initial noise) from sample_id: the engine '
+                f'replays steps after crash recovery and relies on identical results.')
+            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):
+            global _called_from_inside_tl_decorator
+            _called_from_inside_tl_decorator += 1
+            try:
+                result = user_function(sample_id, prev_inputs, prev_outputs, preprocess_response)
+            finally:
+                _called_from_inside_tl_decorator -= 1
+            return result
+
+        leap_binder.set_autoregressive_step(inner_without_validate)
+
+        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)
+
+            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:
+                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, 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,
+                                                       preprocess_response)
+                _assert_deterministic(result, second_result)
+
+            _validate_result(result, is_first_call=prev_inputs is None)
+            _validate_consistency(result)
+
+            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 result
+
+        class _MappingInputsPlaceholder:
+            """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 non-passthrough key, once.
+            """
+            def __init__(self):
+                self._items: Dict[str, Any] = {}
+
+            def __getitem__(self, key):
+                assert isinstance(key, str), \
+                    f'Expected a string model input name, got {type(key)} instead.'
+                if key not in self._items:
+                    class TempMapping:
+                        pass
+                    ret = TempMapping()
+                    ret.node_mapping = NodeMapping(key, NodeMappingType.Input)
+                    if not key.startswith('_'):
+                        leap_binder.mapping_connections.append(NodeConnection(ret.node_mapping, None))
+                    self._items[key] = ret
+                return self._items[key]
+
+            def get(self, key, default=None):
+                return self[key]
+
+            def _unsupported_iteration(self):
+                raise Exception(
+                    'Iterating over the autoregressive hook result is not supported inside the '
+                    'integration test. Pass the model its inputs with explicit keys, e.g. '
+                    "model.run(None, {'x': inputs['x'], 't': inputs['t']}).")
+
+            def items(self):
+                self._unsupported_iteration()
+
+            def keys(self):
+                self._unsupported_iteration()
+
+            def values(self):
+                self._unsupported_iteration()
+
+        mapping_placeholder = _MappingInputsPlaceholder()
+
+        def mapping_inner(*args, **kwargs):
+            return mapping_placeholder
+
+        def final_inner(*args, **kwargs):
+            if os.environ.get(mapping_runtime_mode_env_var_mame):
+                return mapping_inner(*args, **kwargs)
+            else:
+                return inner(*args, **kwargs)
+
+        if not _call_from_tl_platform:
+            register_decorator_row("tensorleap_autoregressive_step")
+
+        return final_inner
+
+    return decorating_function
+
+
 def tensorleap_preprocess():
     def decorating_function(user_function: Callable[[], List[PreprocessResponse]]):
         leap_binder.set_preprocess(user_function)
@@ -1929,8 +2213,13 @@ def tensorleap_input_encoder(name: str, channel_dim=_UNSET, model_input_index=No
 
                 )
 
+        # channel_dim is the channel axis in the BATCHED tensor (batch = axis 0):
+        # 1 for channel-first (NCHW), -1 for channel-last (NHWC); 0 is invalid.
         if channel_dim <= 0 and channel_dim != -1:
-            raise Exception(f"Channel dim for input {name} is expected to be either -1 or positive")
+            raise Exception(
+                f"Invalid channel_dim ({channel_dim}) for input '{name}'. channel_dim is the "
+                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, \
@@ -2149,6 +2438,7 @@ def tensorleap_custom_loss(name: str, connects_to=None):
             _called_from_inside_tl_decorator += 1
 
             try:
+                _require_sample_preprocess_response_supplied(user_function, args, kwargs)
                 result = user_function(*args, **kwargs)
             finally:
                 _called_from_inside_tl_decorator -= 1
@@ -2394,6 +2684,12 @@ def tensorleap_status_table():
             table.append({"name": "tensorleap simulation (optional)", "Added to integration": UNKNOWN})
         _sim_tracking[sim_name] = "registered"
 
+    def register_decorator_row(name: str):
+        # Rows for decorators that only apply to some integration kinds (e.g. the autoregressive
+        # step hook) are appended on registration, so unrelated integrations never show them.
+        if _find_row(_remove_suffix(name, " (optional)")) is None:
+            table.append({"name": name, "Added to integration": UNKNOWN})
+
     def mark_sim_result(sim_name: str, passed: bool):
         _sim_tracking[sim_name] = "passed" if passed else "failed"
 
@@ -2440,11 +2736,12 @@ def tensorleap_status_table():
     atexit.register(run_on_exit)
     sys.excepthook = handle_exception
 
-    return set_current, update_env_params, register_sim, mark_sim_result
+    return set_current, update_env_params, register_sim, mark_sim_result, register_decorator_row
 
 
 if not _call_from_tl_platform:
-    set_current, update_env_params_func, register_sim, mark_sim_result = tensorleap_status_table()
+    set_current, update_env_params_func, register_sim, mark_sim_result, register_decorator_row = \
+        tensorleap_status_table()
 
 
 

code_loader/leaploader.py

@@ -64,9 +64,6 @@ class LeapLoader(LeapLoaderBase):
             self.evaluate_module()
             if global_leap_binder.integration_test_func is not None:
                 global_leap_binder.integration_test_func(None, PreprocessResponse(state=DataStateType.training, length=0))
-        except DatasetScriptException:
-            global_leap_binder.setup_container = DatasetIntegrationSetup()
-            raise
         except TypeError as e:
             import traceback
             global_leap_binder.setup_container = DatasetIntegrationSetup()
@@ -88,17 +85,10 @@ class LeapLoader(LeapLoaderBase):
                 return
 
             parent_path = str(Path(full_path).parent)
-            if parent_path == full_path:
-                return
             append_path_recursively(parent_path)
             sys.path.append(parent_path)
 
         file_path = Path(self.code_path, self.code_entry_name)
-        if not file_path.is_file():
-            raise DatasetScriptException(
-                f"Entry file not found: '{self.code_entry_name}' does not exist under '{self.code_path}'. "
-                f"Check the 'entryFile' setting in your leap.yaml."
-            )
         append_path_recursively(str(file_path))
 
         importlib.invalidate_caches()
@@ -113,13 +103,6 @@ class LeapLoader(LeapLoaderBase):
 
         spec.loader.exec_module(file)
 
-        if global_leap_binder.setup_container.preprocess is None:
-            raise DatasetScriptException(
-                f"No Tensorleap integration found in '{self.code_entry_name}': the script never called "
-                f"leap_binder.set_preprocess(...). Check the 'entryFile' setting in your leap.yaml — "
-                f"it may be pointing to the wrong file."
-            )
-
     @lru_cache()
     def metric_by_name(self) -> Dict[str, MetricHandlerData]:
         self.exec_script()
@@ -337,12 +320,16 @@ class LeapLoader(LeapLoaderBase):
             try:
                 self.exec_script()
 
+                global_leap_binder.validate_autoregressive_setup()
                 preprocess_test_payload = self._check_preprocess()
                 test_payloads.append(preprocess_test_payload)
                 handlers_test_payloads = self._check_handlers()
                 test_payloads.extend(handlers_test_payloads)
                 simulation_test_payloads = self._check_simulations()
                 test_payloads.extend(simulation_test_payloads)
+                autoregressive_test_payload = self._check_autoregressive_step()
+                if autoregressive_test_payload is not None:
+                    test_payloads.append(autoregressive_test_payload)
                 is_valid = all([payload.is_passed for payload in test_payloads])
                 setup_response = self.get_dataset_setup_response(handlers_test_payloads)
 
@@ -462,6 +449,69 @@ class LeapLoader(LeapLoaderBase):
             result_payloads.append(test_result)
         return result_payloads
 
+    def _check_autoregressive_step(self) -> Optional[DatasetTestResultPayload]:
+        handler = global_leap_binder.setup_container.autoregressive_step
+        if handler is None:
+            return None
+        test_result = DatasetTestResultPayload(handler.name)
+        try:
+            preprocess_result = self._preprocess_result()
+            input_shapes: Dict[str, List[int]] = {}
+            for state, preprocess_response in preprocess_result.items():
+                if preprocess_response.sample_ids_to_instance_mappings:
+                    raise Exception('Element instances are not supported together with '
+                                    'tensorleap_autoregressive_step.')
+                sample_id = preprocess_response.sample_ids[0]
+                first_result = handler.function(sample_id, None, None, preprocess_response)
+                second_result = handler.function(sample_id, None, None, preprocess_response)
+                if not isinstance(first_result, dict):
+                    raise Exception('The autoregressive step hook must return a dict of initial model '
+                                    f'inputs on its first call, got {type(first_result)} '
+                                    f'(state: {state.name}).')
+                model_input_items = {key: value for key, value in first_result.items()
+                                     if not key.startswith('_')}
+                if not model_input_items:
+                    raise Exception('The autoregressive step hook returned no model inputs on its first '
+                                    'call — every key starts with "_" (underscore keys are passthrough '
+                                    f'state, never fed to the model). State: {state.name}.')
+                for key, value in model_input_items.items():
+                    if not isinstance(value, np.ndarray):
+                        raise Exception(f'The autoregressive step hook returned a non-numpy value for '
+                                        f'model input "{key}" ({type(value)}). State: {state.name}.')
+                    if key in input_shapes and input_shapes[key] != list(value.shape):
+                        raise Exception(f'The autoregressive step hook returned different shapes for '
+                                        f'model input "{key}" across states: {input_shapes[key]} vs '
+                                        f'{list(value.shape)}. Shapes must be fixed.')
+                    input_shapes[key] = list(value.shape)
+                    if not np.array_equal(value, second_result[key]):
+                        raise Exception('The autoregressive step hook is not deterministic — two calls '
+                                        'with identical arguments returned different outputs for model '
+                                        f'input "{key}". Seed any stochasticity from sample_id: the '
+                                        'engine replays steps after crash recovery and relies on '
+                                        'identical results.')
+            handler.input_shapes = input_shapes
+        except Exception as e:
+            line_number, file_name, stacktrace = get_root_exception_file_and_line_number()
+            test_result.display[TestingSectionEnum.Errors.name] = \
+                f"{repr(e)} in file {file_name}, line_number:  {line_number}\nStacktrace:\n{stacktrace}"
+            test_result.is_passed = False
+        return test_result
+
+    def _ensure_autoregressive_input_shapes(self) -> Dict[str, List[int]]:
+        handler = global_leap_binder.setup_container.autoregressive_step
+        assert handler is not None
+        if handler.input_shapes is None:
+            preprocess_result = self._preprocess_result()
+            first_state_response = next(iter(preprocess_result.values()))
+            first_result = handler.function(first_state_response.sample_ids[0], None, None,
+                                            first_state_response)
+            if not isinstance(first_result, dict):
+                raise Exception('The autoregressive step hook must return a dict of initial model '
+                                f'inputs on its first call, got {type(first_result)}.')
+            handler.input_shapes = {key: list(value.shape) for key, value in first_result.items()
+                                    if not key.startswith('_') and isinstance(value, np.ndarray)}
+        return handler.input_shapes
+
     def run_simulation(self, sim_name, params=None, n_samples=1, seed=0,
                        sample_ids=None, extend_preprocess=True):
         # type: (str, Optional[Dict[str, Any]], int, int, Optional[List[str]], bool) -> Dict[str, Any]
@@ -652,6 +702,12 @@ class LeapLoader(LeapLoaderBase):
                 raise Exception(f"cant calculate shape for input, input name:{inp.name}")
             inputs.append(DatasetInputInstance(name=inp.name, shape=inp.shape, channel_dim=inp.channel_dim))
 
+        if setup.autoregressive_step is not None:
+            # The hook is the sole input source; report its step-0 inputs as the dataset inputs so
+            # graph building, input-coverage validation and type dictionaries see real specs.
+            for input_name, input_shape in self._ensure_autoregressive_input_shapes().items():
+                inputs.append(DatasetInputInstance(name=input_name, shape=input_shape, channel_dim=-1))
+
         ground_truths = []
         for gt in setup.ground_truths:
             if gt.shape is None:
@@ -759,7 +815,21 @@ class LeapLoader(LeapLoaderBase):
         return result_agg
 
     def _get_inputs(self, state: DataStateEnum, sample_id: Union[int, str]) -> Dict[str, npt.NDArray[np.float32]]:
-        return self._get_dataset_handlers(global_leap_binder.setup_container.inputs, state, sample_id)
+        inputs = self._get_dataset_handlers(global_leap_binder.setup_container.inputs, state, sample_id)
+        autoregressive_handler = global_leap_binder.setup_container.autoregressive_step
+        if autoregressive_handler is not None:
+            # Autoregressive integrations have no input encoders — the hook's first call supplies
+            # the sample's (step-0) model inputs, so every sample-fetching flow (graph validation,
+            # analysis, import checks) sees a fully-populated inputs dict.
+            preprocess_response = self._preprocess_result()[state]
+            step_zero_inputs = autoregressive_handler.function(sample_id, None, None, preprocess_response)
+            if not isinstance(step_zero_inputs, dict):
+                raise Exception(
+                    'The autoregressive step hook must return a dict of initial model inputs on its '
+                    f'first call, got {type(step_zero_inputs)} for sample {sample_id}.')
+            inputs.update({key: value for key, value in step_zero_inputs.items()
+                           if not key.startswith('_')})
+        return inputs
 
     def _get_instances_masks(self, state: DataStateEnum, sample_id: Union[int, str], instance_id: int) -> Optional[Dict[str, ElementInstance]]:
         if instance_id is None:
@@ -782,6 +852,12 @@ class LeapLoader(LeapLoaderBase):
     def _get_gt(self, state: DataStateEnum, sample_id: Union[int, str]) -> Dict[str, npt.NDArray[np.float32]]:
         return self._get_dataset_handlers(global_leap_binder.setup_container.ground_truths, state, sample_id)
 
+    def get_gt(self, state: DataStateEnum, sample_id: Union[int, str]) -> Dict[str, npt.NDArray[np.float32]]:
+        """Ground-truth tensors only ({} when no GT encoders are declared) — lets the engine's
+        rollout fetch a chain's GT without re-running the autoregressive hook via get_sample."""
+        self.exec_script()
+        return self._get_gt(state, sample_id)
+
     @lru_cache()
     def _metadata_name_to_type(self) -> Dict[str, DatasetMetadataType]:
         global_leap_binder.check_preprocess(self._preprocess_result())
@@ -877,6 +953,36 @@ class LeapLoader(LeapLoaderBase):
         self.exec_script()
         return global_leap_binder.setup_container.custom_latent_space is not None
 
+    @lru_cache()
+    def has_autoregressive_step(self) -> bool:
+        self.exec_script()
+        return global_leap_binder.setup_container.autoregressive_step is not None
+
+    def run_autoregressive_step(self, sample_id: Union[int, str],
+                                prev_inputs: Optional[Dict[str, npt.NDArray[np.float32]]],
+                                prev_outputs: Optional[Dict[str, npt.NDArray[np.float32]]],
+                                state: DataStateEnum) -> Optional[Dict[str, npt.NDArray[np.float32]]]:
+        """Execute the user's autoregressive step hook against the in-process preprocess result.
+
+        This is the engine rollout's per-step entry point (invoked on the generic pod). The first
+        call per chain passes prev_inputs=None, prev_outputs=None and returns the initial model
+        inputs; a None return ends the chain.
+        """
+        self.exec_script()
+        handler = global_leap_binder.setup_container.autoregressive_step
+        if handler is None:
+            raise Exception('run_autoregressive_step was called but the integration has no '
+                            'tensorleap_autoregressive_step hook.')
+        preprocess_response = self._preprocess_result()[state]
+        return handler.function(sample_id, prev_inputs, prev_outputs, preprocess_response)
+
+    def get_prediction_names_in_order(self) -> List[str]:
+        """Model output names by declared prediction-type order (validated against the model's
+        output count at integration-test time), used to key prev_outputs for the hook."""
+        self.exec_script()
+        return [prediction_type.name for prediction_type in
+                global_leap_binder.setup_container.prediction_types]
+
     def get_instances_data(self, state: DataStateEnum) -> Tuple[Dict[str, List[str]], Dict[str, str]]:
         preprocess_result = self._preprocess_result()
         preprocess_state = preprocess_result.get(state)

code_loader/leaploaderbase.py

@@ -149,6 +149,11 @@ class LeapLoaderBase:
     def has_custom_latent_space_decorator(self) -> bool:
         pass
 
+    def has_autoregressive_step(self) -> bool:
+        # Default False so existing LeapLoaderBase implementations stay valid; the concrete
+        # loaders that can answer (LeapLoader in-process, LeapLoaderWithRedis via RPC) override.
+        return False
+
     @abstractmethod
     def get_heatmap_visualizer_raw_vis_input_arg_name(self, visualizer_name: str) -> Optional[str]:
         pass

{code_loader-1.0.188.dev1.dist-info → code_loader-1.0.189.dev0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: code-loader
-Version: 1.0.188.dev1
+Version: 1.0.189.dev0
 Summary: 
 Home-page: https://github.com/tensorleap/code-loader
 License: MIT

{code_loader-1.0.188.dev1.dist-info → code_loader-1.0.189.dev0.dist-info}/RECORD

@@ -1,7 +1,7 @@
 LICENSE,sha256=qIwWjdspQeSMTtnFZBC8MuT-95L02FPvzRUdWFxrwJY,1067
 code_loader/__init__.py,sha256=outxRQ0M-zMfV0QGVJmAed5qWfRmyD0TV6-goEGAzBw,406
 code_loader/contract/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-code_loader/contract/datasetclasses.py,sha256=RYZcnX7vLMYhyQtCbX2PscVFNAYULxgAA9U4DRRfLqA,10456
+code_loader/contract/datasetclasses.py,sha256=yib7LKHwii0tl0SSMUp1PWgskCs8eI2MQWvtoHs5qgs,11630
 code_loader/contract/enums.py,sha256=2q-IV_5g9lLE306DIbWA1c0tn5IhDtxsKxyV1x_Lreg,1671
 code_loader/contract/exceptions.py,sha256=jWqu5i7t-0IG0jGRsKF4DjJdrsdpJjIYpUkN1F4RiyQ,51
 code_loader/contract/mapping.py,sha256=sWJhpng-IkOzQnWQdMT5w2ZZ3X1Z_OOzSwCLXIS7oxE,1446
@@ -21,10 +21,10 @@ code_loader/experiment_api/types.py,sha256=MY8xFARHwdVA7p4dxyhD60ShmttgTvb4qdp1o
 code_loader/experiment_api/utils.py,sha256=XZHtxge12TS4H4-8PjV3sKuhp8Ud6ojAiIzTZJEqBqc,3304
 code_loader/experiment_api/workingspace_config_utils.py,sha256=DLzXQCg4dgTV_YgaSbeTVzq-2ja_SQw4zi7LXwKL9cY,990
 code_loader/inner_leap_binder/__init__.py,sha256=koOlJyMNYzGbEsoIbXathSmQ-L38N_pEXH_HvL7beXU,99
-code_loader/inner_leap_binder/leapbinder.py,sha256=XLYYcV50qjMvoC1S6WW0tLBch_0g5gl1UyHiVSWYbvg,40491
-code_loader/inner_leap_binder/leapbinder_decorators.py,sha256=qfKa--0kG_K_Eb8wRd5tIXEJ6rbKva3360gU8EU9eKo,118375
-code_loader/leaploader.py,sha256=ZHfYzGpmb_XrZbzBkn12Mh0CE4tDFcFK-PdrfrmJjHo,45536
-code_loader/leaploaderbase.py,sha256=l36qDA00GhZEG5NLKpEtAXgWJA-UQQIhNFGxywK7mUA,6530
+code_loader/inner_leap_binder/leapbinder.py,sha256=6wnjzTzMgWtypCNyIwubMSYRfY4eZOhSLyNEYZf2fvg,45680
+code_loader/inner_leap_binder/leapbinder_decorators.py,sha256=Z5TBE8HnljJCwY1KY43p1zRqd7zh_eF7nxh2oYfdPDE,134547
+code_loader/leaploader.py,sha256=_Ylgu8vTJoy30_ZDymPZ9pZiSLPbIn85dgC8HVKhzbU,52983
+code_loader/leaploaderbase.py,sha256=2UPunTcEZRs99cbFXkNDWdFrs5nQ7EG_T2qwZBepEXY,6788
 code_loader/mixpanel_tracker.py,sha256=rNwRmFifNbdUoqLQvvhhgpKczWpWiEmd8MfyJe27sxw,9131
 code_loader/plot_functions/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 code_loader/plot_functions/plot_functions.py,sha256=2DC-zlVaN13P4VNx5d8csgs80C6SisaeP1-Kq2LW7iM,16075
@@ -32,7 +32,7 @@ code_loader/plot_functions/visualize.py,sha256=gsBAYYkwMh7jIpJeDMPS8G4CW-pxwx6Lz
 code_loader/utils.py,sha256=YecipkdTA-VcE9F0RQcY9cFnY8P3AksPnHM2Db7xUSk,3972
 code_loader/visualizers/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 code_loader/visualizers/default_visualizers.py,sha256=onRnLE_TXfgLN4o52hQIOOhUcFexGlqJ3xSpQDVLuZM,2604
-code_loader-1.0.188.dev1.dist-info/LICENSE,sha256=qIwWjdspQeSMTtnFZBC8MuT-95L02FPvzRUdWFxrwJY,1067
-code_loader-1.0.188.dev1.dist-info/METADATA,sha256=7iylr6wh5362UPoXBVadO3qrOw-lXaaotywDkMoJqoY,1095
-code_loader-1.0.188.dev1.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-code_loader-1.0.188.dev1.dist-info/RECORD,,
+code_loader-1.0.189.dev0.dist-info/LICENSE,sha256=qIwWjdspQeSMTtnFZBC8MuT-95L02FPvzRUdWFxrwJY,1067
+code_loader-1.0.189.dev0.dist-info/METADATA,sha256=9PckzREIEk-hdCCTiMj-Lw5fzub_mPz2oBxEbzDRSzo,1095
+code_loader-1.0.189.dev0.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+code_loader-1.0.189.dev0.dist-info/RECORD,,