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

code_loader/contract/datasetclasses.py

@@ -1,6 +1,6 @@
 import warnings
 from dataclasses import dataclass, field
-from typing import Any, Callable, List, Optional, Dict, Union, Type, Literal, Tuple
+from typing import Any, Callable, List, Optional, Dict, Union, Type, Literal, cast
 import re
 import numpy as np
 import numpy.typing as npt
@@ -15,6 +15,8 @@ custom_latent_space_attribute = "custom_latent_space"
 
 _simulation_context: Dict[str, bool] = {"active": False}
 
+SampleId = Union[int, str]
+
 
 @dataclass
 class PreprocessResponse:
@@ -39,27 +41,44 @@ class PreprocessResponse:
     """
     length: Optional[int] = None  # Deprecated. Please use sample_ids instead
     data: Any = None
-    sample_ids: Optional[Union[List[str], List[int]]] = None
+    sample_ids: Optional[Union[List[SampleId], List[List[SampleId]]]] = None
     state: Optional[DataStateType] = None
     sample_id_type: Optional[Union[Type[str], Type[int]]] = None
     sample_ids_to_instance_mappings: Optional[Dict[str, List[str]]] = None # in use only for element instance
     instance_to_sample_ids_mappings: Optional[Dict[str, str]] = None # in use only for element instance
     tl_generated: bool = False
+    _grouped: bool = field(default=False, init=False, repr=False, compare=False)
 
     def __post_init__(self) -> None:
         assert self.sample_ids_to_instance_mappings is None, f"Keep sample_ids_to_instance_mappings None when initializing PreprocessResponse"
         assert self.instance_to_sample_ids_mappings is None, f"Keep instance_to_sample_ids_mappings None when initializing PreprocessResponse"
 
         if self.length is not None and self.sample_ids is None:
-            self.sample_ids = [i for i in range(self.length)]
+            self.sample_ids = cast(List[SampleId], [i for i in range(self.length)])
             self.sample_id_type = int
+            self._grouped = False
         elif self.length is None and self.sample_ids is not None:
-            self.length = len(self.sample_ids)
-            if self.sample_id_type is None:
-                self.sample_id_type = str
-            if self.sample_id_type == str:
-                for sample_id in self.sample_ids:
-                    assert isinstance(sample_id, str), f"Sample id should be of type str. Got: {type(sample_id)}"
+            self._grouped = len(self.sample_ids) > 0 and isinstance(self.sample_ids[0], list)
+            if self._grouped:
+                groups = cast(List[List[SampleId]], self.sample_ids)
+                assert len(groups) >= 1, "Grouped PreprocessResponse must have at least one group."
+                for group in groups:
+                    assert isinstance(group, list) and len(group) > 0, \
+                        "Each group in a grouped PreprocessResponse must be a non-empty list."
+                self.sample_id_type = type(groups[0][0])
+                for group in groups:
+                    for sample_id in group:
+                        assert isinstance(sample_id, self.sample_id_type), \
+                            f"Sample id should be of type {self.sample_id_type.__name__}. Got: {type(sample_id)}"
+                self.length = sum(len(group) for group in groups)
+            else:
+                flat_ids = cast(List[SampleId], self.sample_ids)
+                self.length = len(flat_ids)
+                if self.sample_id_type is None:
+                    self.sample_id_type = str
+                if self.sample_id_type == str:
+                    for sample_id in flat_ids:
+                        assert isinstance(sample_id, str), f"Sample id should be of type str. Got: {type(sample_id)}"
         else:
             raise Exception("length is deprecated, please use sample_ids instead.")
 
@@ -79,9 +98,36 @@ class PreprocessResponse:
         return id(self)
 
     def __len__(self) -> int:
+        assert self.length is not None
+        return self.length
+
+    @property
+    def is_grouped(self) -> bool:
+        return self._grouped
+
+    @property
+    def num_groups(self) -> int:
+        # Grouped: the number of groups. Flat: the (flat) sample count, since sample_ids is the
+        # flat id list and there is one sample per "group".
         assert self.sample_ids is not None
         return len(self.sample_ids)
 
+    @property
+    def groups(self) -> Optional[List[List[SampleId]]]:
+        # The nested list of groups when grouped, otherwise None. Callers should
+        # branch on is_grouped before relying on this.
+        return cast(List[List[SampleId]], self.sample_ids) if self._grouped else None
+
+    @property
+    def flat_sample_ids(self) -> List[SampleId]:
+        # All sample ids in order: concatenation of the groups when grouped,
+        # otherwise the flat sample_ids list as-is.
+        assert self.sample_ids is not None
+        if self._grouped:
+            groups = cast(List[List[SampleId]], self.sample_ids)
+            return [sample_id for group in groups for sample_id in group]
+        return cast(List[SampleId], self.sample_ids)
+
 @dataclass
 class ElementInstance:
     name: str
@@ -92,19 +138,15 @@ 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, state, preprocess) -> (next model inputs | None, state).
-# First call per chain receives prev_inputs=None, prev_outputs=None, state=None and returns the
-# initial model inputs plus the initial state. Returning next_inputs=None ends the chain — the
-# terminating call still returns state, so the final step participates in state aggregation.
-# State is an arbitrary nest of dicts/lists holding numpy arrays, numbers, strings or bools; it is
-# never fed to the model and rides the engine's redis queue on every step, so accumulate
-# reductions, not raw per-step tensors. The hook must be stateless and deterministically seeded
-# (the engine may replay a step after crash recovery and relies on identical results).
-AutoregressiveChainState = Any
+# (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]]],
-     AutoregressiveChainState, PreprocessResponse],
-    Tuple[Optional[Dict[str, npt.NDArray[np.float32]]], AutoregressiveChainState]]
+     PreprocessResponse],
+    Optional[Dict[str, npt.NDArray[np.float32]]]]
 
 MetadataSectionCallableInterface = Union[
     Callable[[Union[int, str], PreprocessResponse], int],
@@ -260,42 +302,16 @@ class CustomLatentSpaceHandler:
 # dominates the model inputs). 'mean': every latent space is the elementwise mean over all steps.
 AUTOREGRESSIVE_LATENT_SPACE_AGGREGATIONS = ('last_step', 'mean')
 
-# Reserved argument names of autoregressive metrics/losses/visualizers, fed implicitly by the
-# platform from the finished chain (per-chain dicts, no batch axis). Any other argument is wired
-# to a ground-truth encoder through the integration test.
-AUTOREGRESSIVE_IMPLICIT_ARG_NAMES = ('inputs', 'outputs', 'state')
-
 
 @dataclass
 class AutoregressiveStepHandler:
     function: AutoregressiveStepCallableInterface
     name: str = 'autoregressive_step'
-    # Per-model-input shapes (unbatched), discovered by running the hook's first call at parse
-    # time. Fills the role InputHandler.shape plays for input encoders.
+    # 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
     latent_space_aggregation: str = 'last_step'
 
-
-# Per-chain, unbatched callables: called once per finished chain with the final step's tensors —
-# inputs/outputs/state are fed implicitly by the platform (dicts, no batch axis; outputs keyed by
-# prediction-type names), remaining args are wired ground-truth encoder results.
-@dataclass
-class AutoregressiveMetricHandler:
-    metric_handler_data: MetricHandlerData
-    function: CustomCallableInterfaceMultiArgs
-
-
-@dataclass
-class AutoregressiveLossHandler:
-    custom_loss_handler_data: CustomLossHandlerData
-    function: CustomCallableInterface
-
-
-@dataclass
-class AutoregressiveVisualizerHandler:
-    visualizer_handler_data: VisualizerHandlerData
-    function: VisualizerCallableInterface
-
 @dataclass
 class PredictionTypeHandler:
     name: str
@@ -337,9 +353,6 @@ class DatasetIntegrationSetup:
     custom_latent_space: Optional[CustomLatentSpaceHandler] = None
     simulations: List[SimulationHandler] = field(default_factory=list)
     autoregressive_step: Optional[AutoregressiveStepHandler] = None
-    autoregressive_metrics: List[AutoregressiveMetricHandler] = field(default_factory=list)
-    autoregressive_losses: List[AutoregressiveLossHandler] = field(default_factory=list)
-    autoregressive_visualizers: List[AutoregressiveVisualizerHandler] = field(default_factory=list)
 
 
 @dataclass

code_loader/inner_leap_binder/leapbinder.py

@@ -1,6 +1,6 @@
 
 import inspect
-from typing import Callable, List, Optional, Dict, Any, Type, Union, get_args
+from typing import Callable, List, Optional, Dict, Any, Type, Union, get_args, cast
 
 import numpy as np
 import numpy.typing as npt
@@ -14,8 +14,7 @@ from code_loader.contract.datasetclasses import SectionCallableInterface, InputH
     RawInputsForHeatmap, VisualizerHandlerData, MetricHandlerData, CustomLossHandlerData, SamplePreprocessResponse, \
     ElementInstanceMasksHandler, InstanceCallableInterface, CustomLatentSpaceHandler, InstanceMetricHandler, \
     SimulationHandler, _simulation_context, AutoregressiveStepHandler, AutoregressiveStepCallableInterface, \
-    AUTOREGRESSIVE_LATENT_SPACE_AGGREGATIONS, AUTOREGRESSIVE_IMPLICIT_ARG_NAMES, \
-    AutoregressiveMetricHandler, AutoregressiveLossHandler, AutoregressiveVisualizerHandler
+    AUTOREGRESSIVE_LATENT_SPACE_AGGREGATIONS
 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
@@ -554,83 +553,6 @@ class LeapBinder:
 
         self.set_metadata(builtin_chain_metadata, 'builtin_chain')
 
-    def _autoregressive_arg_names(self, function: Callable[..., Any], decorator_name: str,
-                                  name: str) -> List[str]:
-        """Split an autoregressive callable's signature into implicit chain args (validated
-        present) and wired arg names (returned) — the wired args connect to ground-truth
-        encoders through the integration test."""
-        argspec = inspect.getfullargspec(function)
-        for arg_name, arg_type in argspec.annotations.items():
-            _reject_stringized_sample_preprocess_response(function, arg_name, arg_type)
-            if arg_type == SamplePreprocessResponse:
-                raise Exception(
-                    f'{decorator_name} "{name}": SamplePreprocessResponse arguments are not '
-                    f'supported for autoregressive decorators yet. Derive what you need inside '
-                    f'the autoregressive step hook and carry it in the state.')
-        implicit = [arg_name for arg_name in argspec[0]
-                    if arg_name in AUTOREGRESSIVE_IMPLICIT_ARG_NAMES]
-        if not implicit:
-            raise Exception(
-                f'{decorator_name} "{name}": the function must declare at least one of the '
-                f'implicit chain arguments {list(AUTOREGRESSIVE_IMPLICIT_ARG_NAMES)} — they are '
-                f'fed by the platform from the finished chain.')
-        return [arg_name for arg_name in argspec[0]
-                if arg_name not in AUTOREGRESSIVE_IMPLICIT_ARG_NAMES]
-
-    def _assert_unique_metric_name(self, name: str) -> None:
-        existing = [handler.metric_handler_data.name for handler in self.setup_container.metrics] \
-            + [handler.metric_handler_data.name
-               for handler in self.setup_container.autoregressive_metrics]
-        if name in existing:
-            raise Exception(f'Metric with name {name} already exists. Please choose another')
-
-    def _assert_unique_loss_name(self, name: str) -> None:
-        existing = [handler.custom_loss_handler_data.name
-                    for handler in self.setup_container.custom_loss_handlers] \
-            + [handler.custom_loss_handler_data.name
-               for handler in self.setup_container.autoregressive_losses]
-        if name in existing:
-            raise Exception(f'Custom loss with name {name} already exists. Please choose another')
-
-    def _assert_unique_visualizer_name(self, name: str) -> None:
-        existing = [handler.visualizer_handler_data.name
-                    for handler in self.setup_container.visualizers] \
-            + [handler.visualizer_handler_data.name
-               for handler in self.setup_container.autoregressive_visualizers]
-        if name in existing:
-            raise Exception(f'Visualizer with name {name} already exists. Please choose another')
-
-    def add_autoregressive_metric(self, function: CustomCallableInterfaceMultiArgs, name: str,
-                                  direction: Optional[Union[MetricDirection, Dict[str, MetricDirection]]]
-                                  = MetricDirection.Downward,
-                                  compute_insights: Optional[Union[bool, Dict[str, bool]]] = None) -> None:
-        self._assert_unique_metric_name(name)
-        arg_names = self._autoregressive_arg_names(function, 'tensorleap_autoregressive_metric',
-                                                   name)
-        metric_handler_data = MetricHandlerData(name, arg_names, direction, compute_insights)
-        self.setup_container.autoregressive_metrics.append(
-            AutoregressiveMetricHandler(metric_handler_data, function))
-
-    def add_autoregressive_loss(self, function: CustomCallableInterface, name: str) -> None:
-        self._assert_unique_loss_name(name)
-        arg_names = self._autoregressive_arg_names(function, 'tensorleap_autoregressive_loss',
-                                                   name)
-        self.setup_container.autoregressive_losses.append(
-            AutoregressiveLossHandler(CustomLossHandlerData(name, arg_names), function))
-
-    def add_autoregressive_visualizer(self, function: VisualizerCallableInterface, name: str,
-                                      visualizer_type: LeapDataType) -> None:
-        self._assert_unique_visualizer_name(name)
-        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}, '
-                f'should be one of : {", ".join([arg.__name__ for arg in get_args(LeapData)])}')
-        arg_names = self._autoregressive_arg_names(
-            function, 'tensorleap_autoregressive_visualizer', name)
-        self.setup_container.autoregressive_visualizers.append(
-            AutoregressiveVisualizerHandler(VisualizerHandlerData(name, visualizer_type,
-                                                                  arg_names), function))
-
     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:
         """
@@ -713,6 +635,15 @@ class LeapBinder:
             if state_enum in preprocess_result_dict:
                 preprocess_result_dict_in_correct_order[state_enum] = preprocess_result_dict[state_enum]
 
+        # All-or-none: every state must agree on grouped (nested sample_ids) vs flat.
+        if len({r.is_grouped for r in preprocess_result_dict_in_correct_order.values()}) > 1:
+            shapes = ', '.join(
+                f"{state.name}={'grouped' if r.is_grouped else 'flat'}"
+                for state, r in preprocess_result_dict_in_correct_order.items())
+            raise Exception(
+                "All preprocess states must be either all grouped (nested sample_ids) or all flat, "
+                f"but got a mix: {shapes}.")
+
         return preprocess_result_dict_in_correct_order
 
     def get_preprocess_unlabeled_result(self) -> Optional[PreprocessResponse]:
@@ -733,7 +664,15 @@ class LeapBinder:
             preprocess_response: PreprocessResponse, test_result: List[DatasetTestResultPayload],
             dataset_base_handler: Union[DatasetBaseHandler, MetadataHandler], state: DataStateEnum) -> List[DatasetTestResultPayload]:
         assert preprocess_response.sample_ids is not None
-        raw_result = dataset_base_handler.function(preprocess_response.sample_ids[0], preprocess_response)
+        if preprocess_response.is_grouped:
+            # Grouped: probe with the first group, then reduce to a single sample's result so the
+            # recorded shape/type is per-sample (matching the flat case), not the (B, *) batch.
+            # (casts: sample_ids[0] is a group list here, and the grouped result is a per-sample list.)
+            group = preprocess_response.sample_ids[0]
+            raw_result = cast(Any, dataset_base_handler.function(cast(Any, group), preprocess_response))[0]
+        else:
+            raw_result = dataset_base_handler.function(
+                cast(Union[int, str], preprocess_response.sample_ids[0]), preprocess_response)
         handler_type = 'metadata' if isinstance(dataset_base_handler, MetadataHandler) else None
         if isinstance(dataset_base_handler, MetadataHandler):
             if isinstance(raw_result, dict):
@@ -829,10 +768,11 @@ class LeapBinder:
             preprocess_response.tl_generated = True
             if not preprocess_response.length or preprocess_response.length < 1:
                 raise Exception("Simulation '{}' returned PreprocessResponse with length < 1".format(sim.name))
-            preprocess_response.sample_ids = [0]
+            preprocess_response.sample_ids = cast(List[Union[int, str]], [0])
+            sim_sample_ids = cast(List[Union[int, str]], preprocess_response.sample_ids)
             for handler in self.setup_container.inputs:
-                out1 = handler.function(preprocess_response.sample_ids[0], preprocess_response)
-                out2 = handler.function(preprocess_response.sample_ids[0], preprocess_response)
+                out1 = handler.function(sim_sample_ids[0], preprocess_response)
+                out2 = handler.function(sim_sample_ids[0], preprocess_response)
                 if not np.allclose(out1, out2):
                     raise Exception(
                         "Simulation '{}': encoder '{}' is non-deterministic — consecutive calls with seed=0 returned different outputs".format(
@@ -856,23 +796,6 @@ class LeapBinder:
         appear in any order in the integration file.
         """
         if self.setup_container.autoregressive_step is None:
-            declared_autoregressive = [
-                ('tensorleap_autoregressive_metric',
-                 [handler.metric_handler_data.name
-                  for handler in self.setup_container.autoregressive_metrics]),
-                ('tensorleap_autoregressive_loss',
-                 [handler.custom_loss_handler_data.name
-                  for handler in self.setup_container.autoregressive_losses]),
-                ('tensorleap_autoregressive_visualizer',
-                 [handler.visualizer_handler_data.name
-                  for handler in self.setup_container.autoregressive_visualizers]),
-            ]
-            for decorator_name, names in declared_autoregressive:
-                if names:
-                    raise Exception(
-                        f'{decorator_name} {names} requires a tensorleap_autoregressive_step '
-                        f'hook — autoregressive decorators consume the finished chain the hook '
-                        f'drives.')
             return
         if self.setup_container.inputs:
             raise Exception(