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

code_loader/leaploader.py

@@ -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:
@@ -63,6 +64,9 @@ class LeapLoader(LeapLoaderBase):
             os.environ[mapping_runtime_mode_env_var_mame] = 'TRUE'
             self.evaluate_module()
             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()
                 global_leap_binder.integration_test_func(None, PreprocessResponse(state=DataStateType.training, length=0))
         except TypeError as e:
             import traceback
@@ -246,20 +250,16 @@ class LeapLoader(LeapLoaderBase):
             )
 
         preprocess_result = self._preprocess_result()
-        if state == DataStateEnum.unlabeled and sample_id not in preprocess_result[state].flat_sample_ids:
+        if state == DataStateEnum.unlabeled and sample_id not in preprocess_result[state].sample_ids:
             self._preprocess_result(update_unlabeled_preprocess=True)
 
         metadata, metadata_is_none = self.get_metadata(state, sample_id)
 
         custom_latent_space = None
         if global_leap_binder.setup_container.custom_latent_space is not None:
-            latent_fn = global_leap_binder.setup_container.custom_latent_space.function
-            preprocess_state = preprocess_result[state]
-            if preprocess_state.is_grouped:
-                group_ids, pos = self._locate_group(preprocess_state, sample_id)
-                custom_latent_space = self._to_grouped_array(latent_fn(group_ids, preprocess_state))[pos]
-            else:
-                custom_latent_space = latent_fn(sample_id, preprocess_state)
+            custom_latent_space = global_leap_binder.setup_container.custom_latent_space.function(sample_id,
+                                                                                                  preprocess_result[
+                                                                                                      state])
         instance_mask = self._get_instances_masks(state, sample_id, instance_id)
         sample = DatasetSample(inputs=self._get_inputs(state, sample_id),
                                gt=None if state == DataStateEnum.unlabeled else self._get_gt(state, sample_id),
@@ -326,12 +326,16 @@ class LeapLoader(LeapLoaderBase):
 
                 integration_test_test_payload = self._check_integration_test_exists()
                 test_payloads.append(integration_test_test_payload)
+                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)
 
@@ -386,7 +390,7 @@ class LeapLoader(LeapLoaderBase):
             if self.get_sample_id_type() is str:
                 max_allowed_item_size = np.dtype('<U256').itemsize
                 for state, preprocess_response in preprocess_result.items():
-                    sample_ids_array = np.array(preprocess_response.flat_sample_ids)
+                    sample_ids_array = np.array(preprocess_response.sample_ids)
                     if sample_ids_array.dtype.itemsize > max_allowed_item_size:
                         raise Exception(f"Sample id are too long. Max allowed length is 256 charecters.")
 
@@ -468,6 +472,85 @@ 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, 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_inputs)} '
+                                    f'(state: {state.name}).')
+                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.')
+                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}.')
+                    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[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 '
+                                        '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, None,
+                                            first_state_response)
+            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)}.')
+            handler.input_shapes = {key: list(value.shape) for key, value in first_result[0].items()
+                                    if 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]
@@ -658,6 +741,14 @@ 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.
+            # channel_dim=-1 is intentional for AR inputs — there is no InputHandler to carry a
+            # user-declared value; expose it on the decorator if a use case ever needs otherwise.
+            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:
@@ -702,6 +793,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 = {
@@ -747,57 +853,6 @@ class LeapLoader(LeapLoaderBase):
 
         return sample_ids
 
-    @staticmethod
-    def _grouped_dict_keys(rows: List[Dict[str, Any]], handler_name: str) -> List[str]:
-        """Keys for a group of dict-metadata rows, requiring every row to share the same keys.
-        A group must expose a uniform metadata schema; otherwise the per-key lists would silently
-        drop or KeyError on rows that differ from row 0."""
-        keys = list(rows[0].keys())
-        key_set = set(keys)
-        for pos, row in enumerate(rows):
-            if set(row.keys()) != key_set:
-                raise ValueError(
-                    f"Metadata handler {handler_name!r} returned inconsistent dict keys across the "
-                    f"group: row 0 has {sorted(key_set)} but row {pos} has {sorted(row.keys())}. "
-                    f"All rows in a group must share the same metadata keys.")
-        return keys
-
-    @staticmethod
-    def _to_grouped_array(raw: Any) -> Union[npt.NDArray[np.float32], List[npt.NDArray[np.float32]]]:
-        """Normalize a grouped encoder result. A list of uniformly-shaped per-sample arrays is
-        stacked into a single (B, *) array; a list whose samples differ in shape is kept as a
-        list of per-sample arrays (variable-dim groups); an already-grouped array is returned
-        as-is."""
-        if isinstance(raw, list):
-            shapes = {np.asarray(sample).shape for sample in raw}
-            if len(shapes) == 1:
-                return np.stack(raw, axis=0)
-            return raw
-        return raw
-
-    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))
-        if mapping is None:
-            mapping = {}
-            for group in preprocess_state.groups:
-                for pos, sid in enumerate(group):
-                    if sid in mapping:
-                        raise ValueError(
-                            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
-        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]
-
     def _get_dataset_handlers(self, handlers: Iterable[DatasetBaseHandler],
                               state: DataStateEnum, sample_id: Union[int, str]) -> Dict[str, npt.NDArray[np.float32]]:
         result_agg = {}
@@ -809,72 +864,30 @@ class LeapLoader(LeapLoaderBase):
             sample_id = original_local_id
         else:
             preprocess_state = preprocess_result[state]
-        if preprocess_state.is_grouped:
-            # Grouped dataset: call the encoder once with the whole group, index out the sample.
-            group_ids, pos = self._locate_group(preprocess_state, sample_id)
-            for handler in handlers:
-                grouped = self._to_grouped_array(handler.function(group_ids, preprocess_state))
-                result_agg[handler.name] = grouped[pos]
-            return result_agg
         for handler in handlers:
             handler_result = handler.function(sample_id, preprocess_state)
             handler_name = handler.name
             result_agg[handler_name] = handler_result
         return result_agg
 
-    def get_samples(self, state: DataStateEnum, group_ids: List[Union[int, str]]) -> DatasetSample:
-        """Group-aware fetch: hand the whole group to each encoder in a single call and return a
-        grouped DatasetSample (inputs/gt as (B, *), metadata as per-key lists of length B, index =
-        the group). 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.")
-        inputs = {handler.name: self._to_grouped_array(handler.function(group_ids, preprocess_state))
-                  for handler in global_leap_binder.setup_container.inputs}
-        gt = None
-        if state != DataStateEnum.unlabeled:
-            gt = {handler.name: self._to_grouped_array(handler.function(group_ids, preprocess_state))
-                  for handler in global_leap_binder.setup_container.ground_truths}
-
-        metadata: Dict[str, Any] = {}
-        metadata_is_none: Dict[str, Any] = {}
-        for handler in global_leap_binder.setup_container.metadata:
-            rows = handler.function(group_ids, preprocess_state)  # list of B scalars/dicts
-            if rows and isinstance(rows[0], dict):
-                for key in self._grouped_dict_keys(rows, handler.name):
-                    name = "{}_{}".format(handler.name, key)
-                    converted = [self._convert_metadata_to_correct_type(name, row[key]) for row in rows]
-                    metadata[name] = [v for v, _ in converted]
-                    metadata_is_none[name] = [is_none for _, is_none in converted]
-            else:
-                converted = [self._convert_metadata_to_correct_type(handler.name, row) for row in rows]
-                metadata[handler.name] = [v for v, _ in converted]
-                metadata_is_none[handler.name] = [is_none for _, is_none in converted]
-
-        # custom_latent_space is group-aware like the input/GT encoders: hand it the whole group in a
-        # single call so a file-backed latent fn loads the file once, and normalize to (B, d). This
-        # keeps the mandatory group path non-lossy. instance_masks stay None here: they are a
-        # per-(sample, instance) concern that needs an instance_id the group fetch does not carry.
-        custom_latent_space = None
-        if global_leap_binder.setup_container.custom_latent_space is not None:
-            latent_fn = global_leap_binder.setup_container.custom_latent_space.function
-            custom_latent_space = self._to_grouped_array(latent_fn(group_ids, preprocess_state))
-
-        return DatasetSample(inputs=inputs, gt=gt, metadata=metadata, metadata_is_none=metadata_is_none,
-                             index=list(group_ids), state=state, custom_latent_space=custom_latent_space,
-                             instance_masks=None)
-
     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_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 (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}.')
+            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]]:
         if instance_id is None:
@@ -897,6 +910,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())
@@ -953,18 +972,12 @@ class LeapLoader(LeapLoaderBase):
             sample_id = original_local_id
         else:
             preprocess_state = preprocess_result[state]
-        group_ids, pos = (None, None)
-        if preprocess_state.is_grouped:
-            group_ids, pos = self._locate_group(preprocess_state, sample_id)
         for handler in global_leap_binder.setup_container.metadata:
             if requested_metadata_names:
                 if not is_metadata_name_starts_with_handler_name(handler):
                     continue
 
-            if preprocess_state.is_grouped:
-                handler_result = handler.function(group_ids, preprocess_state)[pos]
-            else:
-                handler_result = handler.function(sample_id, preprocess_state)
+            handler_result = handler.function(sample_id, preprocess_state)
             if isinstance(handler_result, dict):
                 for single_metadata_name, single_metadata_result in handler_result.items():
                     handler_name = f'{handler.name}_{single_metadata_name}'
@@ -983,66 +996,6 @@ class LeapLoader(LeapLoaderBase):
 
         return result_agg, is_none
 
-    def _grouped_metadata_for_group(self, preprocess_state: PreprocessResponse,
-                                    group_ids: List[Union[int, str]],
-                                    requested_metadata_names: Optional[List[str]]
-                                    ) -> Tuple[Dict[str, List[Any]], Dict[str, List[bool]]]:
-        def is_wanted(_handler):
-            if not requested_metadata_names:
-                return True
-            for metadata_name in requested_metadata_names:
-                if metadata_name.startswith(_handler.name + '_') or metadata_name == _handler.name:
-                    return True
-            return False
-
-        results: Dict[str, List[Any]] = {}
-        is_none: Dict[str, List[bool]] = {}
-        for handler in global_leap_binder.setup_container.metadata:
-            if requested_metadata_names and not is_wanted(handler):
-                continue
-            rows = handler.function(group_ids, preprocess_state)
-            if rows and isinstance(rows[0], dict):
-                for key in self._grouped_dict_keys(rows, handler.name):
-                    name = f'{handler.name}_{key}'
-                    if requested_metadata_names and name not in requested_metadata_names:
-                        continue
-                    converted = [self._convert_metadata_to_correct_type(name, row[key]) for row in rows]
-                    results[name] = [v for v, _ in converted]
-                    is_none[name] = [n for _, n in converted]
-            else:
-                name = handler.name
-                if requested_metadata_names and name not in requested_metadata_names:
-                    continue
-                converted = [self._convert_metadata_to_correct_type(name, row) for row in rows]
-                results[name] = [v for v, _ in converted]
-                is_none[name] = [n for _, n in converted]
-        return results, is_none
-
-    def get_metadata_multiple_samples(self, state: DataStateEnum, sample_ids: Union[List[int], List[str]],
-                                      requested_metadata_names: Optional[List[str]] = None
-                                      ) -> Tuple[Dict[str, Union[List[str], List[int], List[bool], List[float]]],
-                                                 Dict[str, List[bool]]]:
-        preprocess_state = self._preprocess_result().get(state)
-        if preprocess_state is None or not preprocess_state.is_grouped:
-            return super().get_metadata_multiple_samples(state, sample_ids, requested_metadata_names)
-
-        sample_id_type = self.get_sample_id_type()
-        aggregated_results: Dict[str, List[Any]] = {}
-        aggregated_is_none: Dict[str, List[bool]] = {}
-        group_cache: Dict[int, Tuple[Dict[str, List[Any]], Dict[str, List[bool]]]] = {}
-        for sample_id in sample_ids:
-            sample_id = sample_id_type(sample_id)
-            group_ids, pos = self._locate_group(preprocess_state, sample_id)
-            key = id(group_ids)
-            if key not in group_cache:
-                group_cache[key] = self._grouped_metadata_for_group(
-                    preprocess_state, group_ids, requested_metadata_names)
-            group_results, group_is_none = group_cache[key]
-            for name, values in group_results.items():
-                aggregated_results.setdefault(name, []).append(values[pos])
-                aggregated_is_none.setdefault(name, []).append(group_is_none[name][pos])
-        return aggregated_results, aggregated_is_none
-
     @lru_cache()
     def get_sample_id_type(self) -> Type:
         preprocess_results = list(self._preprocess_result().values())
@@ -1058,6 +1011,152 @@ 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
+
+    @lru_cache()
+    def get_autoregressive_latent_space_aggregation(self) -> str:
+        self.exec_script()
+        handler = global_leap_binder.setup_container.autoregressive_step
+        if handler is None:
+            return 'last_step'
+        return handler.latent_space_aggregation
+
+    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]:
+        """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, 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
+        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]
+        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:
+            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(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()
+        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(next_inputs)} for sample {sample_id}.')
+        for key, value in next_inputs.items():
+            if isinstance(value, np.ndarray) and 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 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
+        hook. AR integrations must declare prediction types (validate_autoregressive_setup) and
+        the count is checked against the model's outputs at integration-test time; the declared
+        ORDER matching the model's output order is the user's responsibility — it is not
+        validated."""
+        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/utils.py

@@ -14,11 +14,6 @@ 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):
-            if len({sample.shape for sample in result}) > 1:
-                # Variable per-sample shapes within a group: keep the per-sample arrays as a list
-                # instead of forcing a stack that would raise an inhomogeneous-shape error.
-                return result
         numpy_result: npt.NDArray[np.float32] = np.array(result)
         return numpy_result
 
@@ -106,3 +101,39 @@ 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
 
+
+
+AUTOREGRESSIVE_STATE_LEAF_TYPES = (np.ndarray, np.generic, int, float, str, bool, type(None))
+
+
+def validate_autoregressive_state_types(value: Any, path: str = 'state') -> None:
+    if isinstance(value, dict):
+        for key, item in value.items():
+            if not isinstance(key, str):
+                raise AssertionError(
+                    f'autoregressive state validation failed: dict key {key!r} at {path} must be '
+                    f'a string. Got {type(key).__name__}.')
+            validate_autoregressive_state_types(item, f'{path}[{key!r}]')
+        return
+    if isinstance(value, (list, tuple)):
+        for i, item in enumerate(value):
+            validate_autoregressive_state_types(item, f'{path}[{i}]')
+        return
+    if not isinstance(value, AUTOREGRESSIVE_STATE_LEAF_TYPES):
+        raise AssertionError(
+            f'autoregressive state validation failed: {path} holds a {type(value).__name__}. '
+            f'State must be built only from numpy arrays, numbers, strings, bools and nests of '
+            f'dicts/lists of those — it is serialized to the platform queue on every chain step.')
+
+
+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):
+        return isinstance(a, np.ndarray) and isinstance(b, np.ndarray) and \
+            a.shape == b.shape and bool(np.array_equal(a, b))
+    return bool(a == b)