code-loader 1.0.195.dev1__py3-none-any.whl → 1.0.196.dev1__py3-none-any.whl

code_loader/leaploader.py

@@ -7,7 +7,7 @@ import sys
 from contextlib import redirect_stdout
 from functools import lru_cache
 from pathlib import Path
-from typing import Dict, List, Iterable, Set, Union, Any, Type, Optional, Callable, Tuple
+from typing import Dict, List, Iterable, Set, FrozenSet, Union, Any, Type, Optional, Callable, Tuple
 
 import numpy as np
 import numpy.typing as npt
@@ -27,7 +27,8 @@ from code_loader.contract.sim_config import FloatBounds, IntBounds, CategoricalB
 from code_loader.inner_leap_binder import global_leap_binder
 from code_loader.inner_leap_binder.leapbinder import mapping_runtime_mode_env_var_mame
 from code_loader.leaploaderbase import LeapLoaderBase
-from code_loader.utils import get_root_exception_file_and_line_number, get_metadata_type_from_variable
+from code_loader.utils import get_root_exception_file_and_line_number, get_metadata_type_from_variable, \
+    validate_autoregressive_state_types, autoregressive_nests_equal
 
 
 def _serialize_sim_bounds(bounds) -> dict:
@@ -76,13 +77,14 @@ class LeapLoader(LeapLoaderBase):
                     preprocess_result = global_leap_binder.get_preprocess_result()
                     self._preprocess_result_cached = preprocess_result
                     is_grouped = any(r.is_grouped for r in preprocess_result.values())
-                except Exception:
-                    is_grouped = False
                 finally:
                     os.environ[mapping_runtime_mode_env_var_mame] = 'TRUE'
 
             _leap_dec._mapping_dataset_is_grouped = is_grouped
             if global_leap_binder.integration_test_func is not None:
+                from code_loader.inner_leap_binder.leapbinder_decorators import \
+                    _reset_model_loop_state
+                _reset_model_loop_state()
                 mapping_preprocess = (
                     PreprocessResponse(sample_ids=[["__mapping_placeholder__"]], state=DataStateType.training)
                     if is_grouped else
@@ -230,7 +232,7 @@ class LeapLoader(LeapLoaderBase):
         additional = self._preprocess_result().get(DataStateEnum.additional)
         if additional is None:
             return set()
-        return set(additional.sample_ids)
+        return set(additional.flat_sample_ids)
 
     def _resolve_synthetic(self, sample_id: Union[int, str],
                            state: Optional[DataStateEnum] = None
@@ -505,31 +507,52 @@ class LeapLoader(LeapLoaderBase):
         try:
             preprocess_result = self._preprocess_result()
             input_shapes: Dict[str, List[int]] = {}
+            expected_keys: Optional[FrozenSet[str]] = None
+            expected_keys_state = ''
             for state, preprocess_response in preprocess_result.items():
                 if preprocess_response.sample_ids_to_instance_mappings:
                     raise Exception('Element instances are not supported together with '
                                     'tensorleap_autoregressive_step.')
-                sample_id = preprocess_response.sample_ids[0]
-                first_result = handler.function(sample_id, None, None, preprocess_response)
-                second_result = handler.function(sample_id, None, None, preprocess_response)
-                if not isinstance(first_result, dict):
+                # Grouped preprocess is fine: grouping only amortizes the shared base reads
+                # (GT/metadata/custom_latent) — the AR hook always drives a single scalar chain
+                # and is never handed a group.
+                sample_id = preprocess_response.flat_sample_ids[0]
+                first_result = handler.function(sample_id, None, None, None, preprocess_response)
+                second_result = handler.function(sample_id, None, None, None, preprocess_response)
+                if not isinstance(first_result, tuple) or len(first_result) != 2:
+                    raise Exception('The autoregressive step hook must return a (next_inputs, '
+                                    f'state) tuple, got {type(first_result)} '
+                                    f'(state: {state.name}).')
+                first_inputs, first_state = first_result
+                if not isinstance(first_inputs, dict) or not first_inputs:
                     raise Exception('The autoregressive step hook must return a dict of initial model '
-                                    f'inputs on its first call, got {type(first_result)} '
+                                    f'inputs on its first call, got {type(first_inputs)} '
                                     f'(state: {state.name}).')
-                if not isinstance(second_result, dict) or \
-                        set(second_result.keys()) != set(first_result.keys()):
+                validate_autoregressive_state_types(first_state)
+                if not isinstance(second_result, tuple) or len(second_result) != 2 or \
+                        not isinstance(second_result[0], dict) or \
+                        set(second_result[0].keys()) != set(first_inputs.keys()) or \
+                        not autoregressive_nests_equal(first_state, second_result[1]):
                     raise Exception('The autoregressive step hook is not deterministic — two calls '
                                     'with identical arguments returned different results '
                                     f'(state: {state.name}). Seed any stochasticity from '
                                     'sample_id: the engine replays steps after crash recovery '
                                     'and relies on identical results.')
-                model_input_items = {key: value for key, value in first_result.items()
-                                     if not key.startswith('_')}
-                if not model_input_items:
-                    raise Exception('The autoregressive step hook returned no model inputs on its first '
-                                    'call — every key starts with "_" (underscore keys are passthrough '
-                                    f'state, never fed to the model). State: {state.name}.')
-                for key, value in model_input_items.items():
+                if expected_keys is None:
+                    expected_keys = frozenset(first_inputs)
+                    expected_keys_state = state.name
+                elif frozenset(first_inputs) != expected_keys:
+                    raise Exception(f'The autoregressive step hook returned different model input '
+                                    f'keys across states: {sorted(expected_keys)} '
+                                    f'({expected_keys_state}) vs {sorted(first_inputs)} '
+                                    f'({state.name}). The key set must be identical for every '
+                                    f'state.')
+                for key, value in first_inputs.items():
+                    if key.startswith('_'):
+                        raise Exception(f'The autoregressive step hook returned model input key '
+                                        f'"{key}" — underscore passthrough keys were replaced by '
+                                        f'the explicit state channel; move passthrough values '
+                                        f'into the returned state. State: {state.name}.')
                     if not isinstance(value, np.ndarray):
                         raise Exception(f'The autoregressive step hook returned a non-numpy value for '
                                         f'model input "{key}" ({type(value)}). State: {state.name}.')
@@ -538,7 +561,7 @@ class LeapLoader(LeapLoaderBase):
                                         f'model input "{key}" across states: {input_shapes[key]} vs '
                                         f'{list(value.shape)}. Shapes must be fixed.')
                     input_shapes[key] = list(value.shape)
-                    if not np.array_equal(value, second_result[key]):
+                    if not autoregressive_nests_equal(value, second_result[0][key]):
                         raise Exception('The autoregressive step hook is not deterministic — two calls '
                                         'with identical arguments returned different outputs for model '
                                         f'input "{key}". Seed any stochasticity from sample_id: the '
@@ -552,19 +575,39 @@ class LeapLoader(LeapLoaderBase):
             test_result.is_passed = False
         return test_result
 
+    @staticmethod
+    def _validate_autoregressive_step_inputs(step_inputs: Dict[str, Any],
+                                             sample_id: Union[int, str]) -> None:
+        # The parse-time check only covers the first sample of each state; the engine runtime
+        # paths must enforce the same key/value contract for every sample they touch.
+        for key, value in step_inputs.items():
+            if key.startswith('_'):
+                raise Exception(f'The autoregressive step hook returned model input key "{key}" '
+                                f'for sample {sample_id} — underscore passthrough keys were '
+                                'replaced by the explicit state channel; move passthrough values '
+                                'into the returned state.')
+            if not isinstance(value, np.ndarray):
+                raise Exception('The autoregressive step hook returned a non-numpy value for '
+                                f'model input "{key}" ({type(value).__name__}) for sample '
+                                f'{sample_id}.')
+
     def _ensure_autoregressive_input_shapes(self) -> Dict[str, List[int]]:
         handler = global_leap_binder.setup_container.autoregressive_step
         assert handler is not None
         if handler.input_shapes is None:
             preprocess_result = self._preprocess_result()
             first_state_response = next(iter(preprocess_result.values()))
-            first_result = handler.function(first_state_response.sample_ids[0], None, None,
+            first_sample_id = first_state_response.flat_sample_ids[0]
+            first_result = handler.function(first_sample_id, None, None, None,
                                             first_state_response)
-            if not isinstance(first_result, dict):
-                raise Exception('The autoregressive step hook must return a dict of initial model '
-                                f'inputs on its first call, got {type(first_result)}.')
-            handler.input_shapes = {key: list(value.shape) for key, value in first_result.items()
-                                    if not key.startswith('_') and isinstance(value, np.ndarray)}
+            if not isinstance(first_result, tuple) or len(first_result) != 2 or \
+                    not isinstance(first_result[0], dict):
+                raise Exception('The autoregressive step hook must return a (next_inputs, state) '
+                                'tuple with a dict of initial model inputs on its first call, '
+                                f'got {type(first_result)}.')
+            self._validate_autoregressive_step_inputs(first_result[0], first_sample_id)
+            handler.input_shapes = {key: list(value.shape)
+                                    for key, value in first_result[0].items()}
         return handler.input_shapes
 
     def run_simulation(self, sim_name, params=None, n_samples=1, seed=0,
@@ -833,6 +876,21 @@ class LeapLoader(LeapLoaderBase):
             metric_inst = MetricInstance(metric.metric_handler_data.name, metric.metric_handler_data.arg_names)
             metrics.append(metric_inst)
 
+        # Autoregressive decorators are reported alongside the regular ones so the platform sees
+        # their names; the engine tells them apart by name via get_autoregressive_decorator_names.
+        for autoregressive_metric in setup.autoregressive_metrics:
+            metrics.append(MetricInstance(autoregressive_metric.metric_handler_data.name,
+                                          autoregressive_metric.metric_handler_data.arg_names))
+        for autoregressive_loss in setup.autoregressive_losses:
+            custom_losses.append(CustomLossInstance(
+                autoregressive_loss.custom_loss_handler_data.name,
+                autoregressive_loss.custom_loss_handler_data.arg_names))
+        for autoregressive_visualizer in setup.autoregressive_visualizers:
+            visualizers.append(VisualizerInstance(
+                autoregressive_visualizer.visualizer_handler_data.name,
+                autoregressive_visualizer.visualizer_handler_data.type,
+                autoregressive_visualizer.visualizer_handler_data.arg_names))
+
         simulations = []
         for sim in setup.simulations:
             sim_config_serialized = {
@@ -909,13 +967,12 @@ class LeapLoader(LeapLoaderBase):
 
     def _locate_group(self, preprocess_state: PreprocessResponse,
                       sample_id: Union[int, str]) -> Tuple[List[Union[int, str]], int]:
-        """Map a sample id to (its group, its position in that group). Cached per
-        PreprocessResponse so grouped single-sample fetches don't rescan the groups."""
-        cache = getattr(self, '_group_pos_cache', None)
-        if cache is None:
-            cache = {}
-            self._group_pos_cache = cache
-        mapping = cache.get(id(preprocess_state))
+        """Map a sample id to (its group, its position in that group). Cached on the
+        PreprocessResponse itself (not keyed by id() on this loader) so grouped single-sample
+        fetches don't rescan the groups, and the cache can never go stale: when the response is
+        replaced (e.g. unlabeled preprocess refresh), the old object and its cache are simply
+        discarded together instead of lingering under a possibly-reused id()."""
+        mapping = preprocess_state._group_pos_cache
         if mapping is None:
             mapping = {}
             for group in preprocess_state.groups:
@@ -925,7 +982,7 @@ class LeapLoader(LeapLoaderBase):
                             f"Duplicate sample id {sid!r} across groups in the preprocess response; "
                             f"sample ids must be unique within and across groups.")
                     mapping[sid] = (group, pos)
-            cache[id(preprocess_state)] = mapping
+            preprocess_state._group_pos_cache = mapping
         if sample_id not in mapping:
             raise KeyError(f"Sample id {sample_id!r} is not present in any group of this preprocess response.")
         return mapping[sample_id]
@@ -961,18 +1018,47 @@ class LeapLoader(LeapLoaderBase):
         samples into the model's B is the engine's job. This is the path a group-aware engine calls
         to read a file once per group. Row order matches group_ids exactly."""
         self.exec_script()
-        preprocess_state = self._preprocess_result()[state]
-        if preprocess_state.is_grouped:
-            # All requested ids must belong to a single group (one file). A cross-group request
-            # would force the encoder to load multiple files, defeating the one-load-per-group
-            # contract; partitioning a scattered request by group is the engine's responsibility.
-            distinct_groups = {id(self._locate_group(preprocess_state, sid)[0]) for sid in group_ids}
-            assert len(distinct_groups) == 1, (
-                f"get_samples received sample ids spanning {len(distinct_groups)} groups; a request "
-                f"must be confined to a single group. The engine partitions cross-group requests and "
-                f"issues one get_samples call per group.")
+        preprocess_result = self._preprocess_result()
+        if state == DataStateEnum.unlabeled and any(
+                sid not in preprocess_result[state].flat_sample_ids for sid in group_ids):
+            # Mirrors get_sample's refresh: the unlabeled preprocess can grow between calls, so a
+            # group_id absent from the current snapshot may just not have been generated yet.
+            self._preprocess_result(update_unlabeled_preprocess=True)
+        preprocess_state = preprocess_result[state]
+        assert preprocess_state.is_grouped, (
+            "get_samples is the group-aware fetch path and requires a grouped preprocess "
+            "response; call get_sample for a flat (non-grouped) dataset.")
+        # All requested ids must belong to a single group (one file). A cross-group request
+        # would force the encoder to load multiple files, defeating the one-load-per-group
+        # contract; partitioning a scattered request by group is the engine's responsibility.
+        distinct_groups = {id(self._locate_group(preprocess_state, sid)[0]) for sid in group_ids}
+        assert len(distinct_groups) == 1, (
+            f"get_samples received sample ids spanning {len(distinct_groups)} groups; a request "
+            f"must be confined to a single group. The engine partitions cross-group requests and "
+            f"issues one get_samples call per group.")
         inputs = {handler.name: self._to_grouped_list(handler.function(group_ids, preprocess_state))
                   for handler in global_leap_binder.setup_container.inputs}
+        autoregressive_handler = global_leap_binder.setup_container.autoregressive_step
+        if autoregressive_handler is not None:
+            # AR integrations have no input encoders — the hook's first call supplies each
+            # chain's step-0 model inputs. The hook is strictly per-scalar-chain (never handed
+            # a group), so it is called once per flat id in the group and the per-sample results
+            # are collected into per-key lists of length B, matching the grouped (never-stacked)
+            # contract used for every other encoder here.
+            step_zero_per_sample = []
+            for sid in group_ids:
+                step_zero_result = autoregressive_handler.function(sid, None, None, None,
+                                                                    preprocess_state)
+                if not isinstance(step_zero_result, tuple) or len(step_zero_result) != 2 or \
+                        not isinstance(step_zero_result[0], dict):
+                    raise Exception(
+                        'The autoregressive step hook must return a (next_inputs, state) tuple '
+                        f'with a dict of initial model inputs on its first call, got '
+                        f'{type(step_zero_result)} for sample {sid}.')
+                self._validate_autoregressive_step_inputs(step_zero_result[0], sid)
+                step_zero_per_sample.append(step_zero_result[0])
+            for key in (step_zero_per_sample[0].keys() if step_zero_per_sample else []):
+                inputs[key] = [row[key] for row in step_zero_per_sample]
         gt = None
         if state != DataStateEnum.unlabeled:
             gt = {handler.name: self._to_grouped_list(handler.function(group_ids, preprocess_state))
@@ -1014,13 +1100,16 @@ class LeapLoader(LeapLoaderBase):
             # the sample's (step-0) model inputs, so every sample-fetching flow (graph validation,
             # analysis, import checks) sees a fully-populated inputs dict.
             preprocess_response = self._preprocess_result()[state]
-            step_zero_inputs = autoregressive_handler.function(sample_id, None, None, preprocess_response)
-            if not isinstance(step_zero_inputs, dict):
+            step_zero_result = autoregressive_handler.function(sample_id, None, None, None,
+                                                               preprocess_response)
+            if not isinstance(step_zero_result, tuple) or len(step_zero_result) != 2 or \
+                    not isinstance(step_zero_result[0], dict):
                 raise Exception(
-                    'The autoregressive step hook must return a dict of initial model inputs on its '
-                    f'first call, got {type(step_zero_inputs)} for sample {sample_id}.')
-            inputs.update({key: value for key, value in step_zero_inputs.items()
-                           if not key.startswith('_')})
+                    'The autoregressive step hook must return a (next_inputs, state) tuple with a '
+                    f'dict of initial model inputs on its first call, got {type(step_zero_result)} '
+                    f'for sample {sample_id}.')
+            self._validate_autoregressive_step_inputs(step_zero_result[0], sample_id)
+            inputs.update(step_zero_result[0])
         return inputs
 
     def _get_instances_masks(self, state: DataStateEnum, sample_id: Union[int, str], instance_id: int) -> Optional[Dict[str, ElementInstance]]:
@@ -1227,12 +1316,15 @@ class LeapLoader(LeapLoaderBase):
     def run_autoregressive_step(self, sample_id: Union[int, str],
                                 prev_inputs: Optional[Dict[str, npt.NDArray[np.float32]]],
                                 prev_outputs: Optional[Dict[str, npt.NDArray[np.float32]]],
-                                state: DataStateEnum) -> Optional[Dict[str, npt.NDArray[np.float32]]]:
+                                chain_state: Any,
+                                state: DataStateEnum
+                                ) -> Tuple[Optional[Dict[str, npt.NDArray[np.float32]]], Any]:
         """Execute the user's autoregressive step hook against the in-process preprocess result.
 
         This is the engine rollout's per-step entry point (invoked on the generic pod). The first
-        call per chain passes prev_inputs=None, prev_outputs=None and returns the initial model
-        inputs; a None return ends the chain.
+        call per chain passes prev_inputs=None, prev_outputs=None, chain_state=None and returns
+        the initial model inputs plus the initial state; a None next_inputs ends the chain — the
+        terminating call still returns the final state.
         """
         self.exec_script()
         handler = global_leap_binder.setup_container.autoregressive_step
@@ -1240,28 +1332,119 @@ class LeapLoader(LeapLoaderBase):
             raise Exception('run_autoregressive_step was called but the integration has no '
                             'tensorleap_autoregressive_step hook.')
         preprocess_response = self._preprocess_result()[state]
-        result = handler.function(sample_id, prev_inputs, prev_outputs, preprocess_response)
-        if result is None:
-            return None
-        if not isinstance(result, dict):
+        result = handler.function(sample_id, prev_inputs, prev_outputs, chain_state,
+                                  preprocess_response)
+        if not isinstance(result, tuple) or len(result) != 2:
+            raise Exception('The autoregressive step hook must return a (next_inputs, state) '
+                            f'tuple, got {type(result)} for sample {sample_id}.')
+        next_inputs, new_chain_state = result
+        validate_autoregressive_state_types(new_chain_state)
+        if next_inputs is None:
+            if prev_inputs is None:
+                raise Exception('The autoregressive step hook returned None for next_inputs on '
+                                f'the first call of the chain (sample {sample_id}) — the first '
+                                'call must return the initial model inputs; None would produce '
+                                'an empty chain with no forward pass.')
+            return None, new_chain_state
+        if not isinstance(next_inputs, dict):
             raise Exception('The autoregressive step hook must return a dict of model inputs or '
-                            f'None to end the chain, got {type(result)} for sample {sample_id}.')
+                            f'None to end the chain, got {type(next_inputs)} for sample '
+                            f'{sample_id}.')
         # Parse-time validation only exercises step 0; enforce the fixed key-set/shape contract
         # on every step here so drift fails with an actionable message instead of a cryptic
         # engine error on a compiled graph.
         expected_shapes = self._ensure_autoregressive_input_shapes()
-        model_inputs = {key: value for key, value in result.items() if not key.startswith('_')}
-        if set(model_inputs) != set(expected_shapes):
+        if set(next_inputs) != set(expected_shapes):
             raise Exception('The autoregressive step hook must return the same model input keys '
                             f'on every call. Expected {sorted(expected_shapes)}, got '
-                            f'{sorted(model_inputs)} for sample {sample_id}.')
-        for key, value in model_inputs.items():
-            if isinstance(value, np.ndarray) and list(value.shape) != expected_shapes[key]:
+                            f'{sorted(next_inputs)} for sample {sample_id}.')
+        for key, value in next_inputs.items():
+            if not isinstance(value, np.ndarray):
+                raise Exception('The autoregressive step hook must return numpy arrays for every '
+                                f'model input on every call — input "{key}" is a '
+                                f'{type(value).__name__} for sample {sample_id} (forgot .numpy() '
+                                'on a fed-back model output?).')
+            if list(value.shape) != expected_shapes[key]:
                 raise Exception('The autoregressive step hook must return fixed tensor shapes on '
                                 f'every call (use fixed-length padding + mask for growing '
                                 f'sequences). Input "{key}" was {expected_shapes[key]} on the '
                                 f'first call and {list(value.shape)} now (sample {sample_id}).')
-        return result
+        return next_inputs, new_chain_state
+
+    @lru_cache()
+    def autoregressive_visualizer_by_name(self) -> Dict[str, VisualizerHandlerData]:
+        self.exec_script()
+        return {handler.visualizer_handler_data.name: handler.visualizer_handler_data
+                for handler in global_leap_binder.setup_container.autoregressive_visualizers}
+
+    @lru_cache()
+    def get_autoregressive_decorator_names(self) -> Dict[str, List[str]]:
+        self.exec_script()
+        setup = global_leap_binder.setup_container
+        return {
+            'metrics': [handler.metric_handler_data.name
+                        for handler in setup.autoregressive_metrics],
+            'losses': [handler.custom_loss_handler_data.name
+                       for handler in setup.autoregressive_losses],
+            'visualizers': [handler.visualizer_handler_data.name
+                            for handler in setup.autoregressive_visualizers],
+        }
+
+    @staticmethod
+    def _autoregressive_callable_kwargs(function: Callable[..., Any],
+                                        inputs: Dict[str, npt.NDArray[np.float32]],
+                                        outputs: Dict[str, npt.NDArray[np.float32]],
+                                        chain_state: Any,
+                                        wired_args: Dict[str, npt.NDArray[np.float32]]
+                                        ) -> Dict[str, Any]:
+        implicit_values = {'inputs': inputs, 'outputs': outputs, 'state': chain_state}
+        kwargs: Dict[str, Any] = {}
+        for arg_name in inspect.getfullargspec(function)[0]:
+            if arg_name in implicit_values:
+                kwargs[arg_name] = implicit_values[arg_name]
+            elif arg_name in wired_args:
+                kwargs[arg_name] = wired_args[arg_name]
+            else:
+                raise Exception(f'Autoregressive callable is missing wired argument '
+                                f'"{arg_name}" — the platform mapping supplies only '
+                                f'{sorted(wired_args)}.')
+        return kwargs
+
+    def run_autoregressive_metric(self, name: str,
+                                  inputs: Dict[str, npt.NDArray[np.float32]],
+                                  outputs: Dict[str, npt.NDArray[np.float32]],
+                                  chain_state: Any,
+                                  wired_args: Dict[str, npt.NDArray[np.float32]]) -> Any:
+        self.exec_script()
+        for handler in global_leap_binder.setup_container.autoregressive_metrics:
+            if handler.metric_handler_data.name == name:
+                return handler.function(**self._autoregressive_callable_kwargs(
+                    handler.function, inputs, outputs, chain_state, wired_args))
+        raise Exception(f'No autoregressive metric named "{name}" is registered.')
+
+    def run_autoregressive_loss(self, name: str,
+                                inputs: Dict[str, npt.NDArray[np.float32]],
+                                outputs: Dict[str, npt.NDArray[np.float32]],
+                                chain_state: Any,
+                                wired_args: Dict[str, npt.NDArray[np.float32]]) -> Any:
+        self.exec_script()
+        for handler in global_leap_binder.setup_container.autoregressive_losses:
+            if handler.custom_loss_handler_data.name == name:
+                return handler.function(**self._autoregressive_callable_kwargs(
+                    handler.function, inputs, outputs, chain_state, wired_args))
+        raise Exception(f'No autoregressive loss named "{name}" is registered.')
+
+    def run_autoregressive_visualizer(self, name: str,
+                                      inputs: Dict[str, npt.NDArray[np.float32]],
+                                      outputs: Dict[str, npt.NDArray[np.float32]],
+                                      chain_state: Any,
+                                      wired_args: Dict[str, npt.NDArray[np.float32]]) -> Any:
+        self.exec_script()
+        for handler in global_leap_binder.setup_container.autoregressive_visualizers:
+            if handler.visualizer_handler_data.name == name:
+                return handler.function(**self._autoregressive_callable_kwargs(
+                    handler.function, inputs, outputs, chain_state, wired_args))
+        raise Exception(f'No autoregressive visualizer named "{name}" is registered.')
 
     def get_prediction_names_in_order(self) -> List[str]:
         """Model output names by declared prediction-type order, used to key prev_outputs for the

code_loader/leaploaderbase.py

@@ -149,10 +149,74 @@ class LeapLoaderBase:
     def has_custom_latent_space_decorator(self) -> bool:
         pass
 
+    # The autoregressive entry points raise instead of `pass`-ing like the abstract methods
+    # above: this class has no ABCMeta, so @abstractmethod is not enforced at runtime and an
+    # un-overridden `pass` body would return None — has_autoregressive_step would silently
+    # report False for an AR integration and the engine would skip the chain rollout with no
+    # error. A stale loader subclass keeps working for non-AR flows and fails loudly here the
+    # moment an AR flow reaches it.
+    @abstractmethod
     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
+        raise NotImplementedError(f'{type(self).__name__} does not implement '
+                                  'has_autoregressive_step — it must be overridden '
+                                  '(LeapLoader in-process, LeapLoaderWithRedis via RPC).')
+
+    @abstractmethod
+    def get_gt(self, state: DataStateEnum, sample_id: Union[int, str]
+               ) -> Dict[str, npt.NDArray[np.float32]]:
+        raise NotImplementedError(f'{type(self).__name__} does not implement get_gt.')
+
+    @abstractmethod
+    def get_prediction_names_in_order(self) -> List[str]:
+        raise NotImplementedError(f'{type(self).__name__} does not implement '
+                                  'get_prediction_names_in_order.')
+
+    @abstractmethod
+    def get_autoregressive_latent_space_aggregation(self) -> str:
+        raise NotImplementedError(f'{type(self).__name__} does not implement '
+                                  'get_autoregressive_latent_space_aggregation.')
+
+    @abstractmethod
+    def get_autoregressive_decorator_names(self) -> Dict[str, List[str]]:
+        raise NotImplementedError(f'{type(self).__name__} does not implement '
+                                  'get_autoregressive_decorator_names.')
+
+    @abstractmethod
+    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]]],
+                                chain_state: Any,
+                                state: DataStateEnum
+                                ) -> Tuple[Optional[Dict[str, npt.NDArray[np.float32]]], Any]:
+        raise NotImplementedError(f'{type(self).__name__} does not implement '
+                                  'run_autoregressive_step.')
+
+    @abstractmethod
+    def run_autoregressive_metric(self, name: str, inputs: Dict[str, npt.NDArray[np.float32]],
+                                  outputs: Dict[str, npt.NDArray[np.float32]], chain_state: Any,
+                                  wired_args: Dict[str, npt.NDArray[np.float32]]) -> Any:
+        raise NotImplementedError(f'{type(self).__name__} does not implement '
+                                  'run_autoregressive_metric.')
+
+    @abstractmethod
+    def run_autoregressive_loss(self, name: str, inputs: Dict[str, npt.NDArray[np.float32]],
+                                outputs: Dict[str, npt.NDArray[np.float32]], chain_state: Any,
+                                wired_args: Dict[str, npt.NDArray[np.float32]]) -> Any:
+        raise NotImplementedError(f'{type(self).__name__} does not implement '
+                                  'run_autoregressive_loss.')
+
+    @abstractmethod
+    def run_autoregressive_visualizer(self, name: str, inputs: Dict[str, npt.NDArray[np.float32]],
+                                      outputs: Dict[str, npt.NDArray[np.float32]],
+                                      chain_state: Any,
+                                      wired_args: Dict[str, npt.NDArray[np.float32]]) -> Any:
+        raise NotImplementedError(f'{type(self).__name__} does not implement '
+                                  'run_autoregressive_visualizer.')
+
+    @abstractmethod
+    def autoregressive_visualizer_by_name(self) -> Dict[str, VisualizerHandlerData]:
+        raise NotImplementedError(f'{type(self).__name__} does not implement '
+                                  'autoregressive_visualizer_by_name.')
 
     @abstractmethod
     def get_heatmap_visualizer_raw_vis_input_arg_name(self, visualizer_name: str) -> Optional[str]:

code_loader/utils.py

@@ -1,3 +1,6 @@
+import io
+import math
+import pickle
 import sys
 from pathlib import Path
 from types import TracebackType
@@ -14,10 +17,12 @@ from code_loader.contract.enums import DatasetMetadataType
 def to_numpy_return_wrapper(encoder_function: SectionCallableInterface) -> SectionCallableInterface:
     def numpy_encoder_function(idx: Union[int, str], samples: PreprocessResponse) -> npt.NDArray[np.float32]:
         result = encoder_function(idx, samples)
-        if isinstance(result, list) and len(result) > 0 and all(isinstance(sample, np.ndarray) for sample in result):
-            # A list of per-sample arrays only comes from a grouped encoder (a flat encoder returns
-            # a single array). Never stack it into a (B, *) batch — grouping is only a storage
-            # read-unit; batching the group to the model's B is the engine's single responsibility.
+        if isinstance(idx, list) and isinstance(result, list) and len(result) > 0 \
+                and all(isinstance(sample, np.ndarray) for sample in result):
+            # Grouped call (idx is a list of sample ids): keep the per-sample list, never stack
+            # it into a (B, *) batch — grouping is only a storage read-unit; batching the group to
+            # the model's B is the engine's single responsibility. Gate on the call shape (idx is a
+            # list), not the result shape, so a flat encoder returning a list of arrays still stacks.
             return result
         numpy_result: npt.NDArray[np.float32] = np.array(result)
         return numpy_result
@@ -106,3 +111,51 @@ def get_metadata_type_from_variable(variable: Union[Optional[str], int, bool, Op
                         f"The return type should be one of [int, float, str, bool]. Got {metadata_type}")
     return dataset_metadata_type
 
+
+
+# The platform queue transports state as pickled payloads which the engine deserializes
+# WITHOUT the integration code loaded, so a class defined in the user's modules round-trips
+# locally but crashes on the engine. Only modules importable on the engine may resolve.
+_STATE_SAFE_PICKLE_MODULES = ('builtins', 'copyreg', 'collections', 'numpy', 'datetime')
+
+
+class _EngineSafeUnpickler(pickle.Unpickler):
+    def find_class(self, module: str, name: str) -> Any:
+        if module.split('.')[0] in _STATE_SAFE_PICKLE_MODULES:
+            return super().find_class(module, name)
+        raise pickle.UnpicklingError(
+            f'{module}.{name} is defined in integration code, which is not importable on the '
+            f'platform engine')
+
+
+def validate_autoregressive_state_types(value: Any, path: str = 'state') -> None:
+    try:
+        _EngineSafeUnpickler(io.BytesIO(pickle.dumps(value))).load()
+    except Exception as e:
+        raise AssertionError(
+            f'autoregressive state validation failed: {path} is not serializable ({e}). State '
+            'is serialized to the platform queue on every chain step and deserialized without '
+            'the integration code — any picklable nest of builtin/numpy values works (dicts, '
+            'lists, tuples, sets, numbers, strings, arrays), but user-defined classes, lambdas '
+            'and open resources do not.') from e
+
+
+def autoregressive_nests_equal(a: Any, b: Any) -> bool:
+    if isinstance(a, dict) or isinstance(b, dict):
+        return isinstance(a, dict) and isinstance(b, dict) and set(a) == set(b) and \
+            all(autoregressive_nests_equal(a[key], b[key]) for key in a)
+    if isinstance(a, (list, tuple)) or isinstance(b, (list, tuple)):
+        return isinstance(a, (list, tuple)) and isinstance(b, (list, tuple)) and \
+            len(a) == len(b) and all(autoregressive_nests_equal(x, y) for x, y in zip(a, b))
+    if isinstance(a, np.ndarray) or isinstance(b, np.ndarray):
+        if not (isinstance(a, np.ndarray) and isinstance(b, np.ndarray) and a.shape == b.shape):
+            return False
+        # NaN leaves must equal themselves: this comparison backs the determinism checks, where
+        # NaN padding / running stats in a perfectly deterministic hook are legitimate.
+        if a.dtype.kind in 'fc' and b.dtype.kind in 'fc':
+            return bool(np.array_equal(a, b, equal_nan=True))
+        return bool(np.array_equal(a, b))
+    if isinstance(a, (float, np.floating)) and isinstance(b, (float, np.floating)) and \
+            math.isnan(a) and math.isnan(b):
+        return True
+    return bool(a == b)

{code_loader-1.0.195.dev1.dist-info → code_loader-1.0.196.dev1.dist-info}/METADATA

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