code-loader 1.0.196.dev1__py3-none-any.whl → 1.0.197.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, cast
+from typing import Any, Callable, List, Optional, Dict, Union, Type, Literal, Tuple
 import re
 import numpy as np
 import numpy.typing as npt
@@ -15,8 +15,6 @@ custom_latent_space_attribute = "custom_latent_space"
 
 _simulation_context: Dict[str, bool] = {"active": False}
 
-SampleId = Union[int, str]
-
 
 @dataclass
 class PreprocessResponse:
@@ -41,46 +39,27 @@ class PreprocessResponse:
     """
     length: Optional[int] = None  # Deprecated. Please use sample_ids instead
     data: Any = None
-    sample_ids: Optional[Union[List[SampleId], List[List[SampleId]]]] = None
+    sample_ids: Optional[Union[List[str], List[int]]] = 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)
-    _group_pos_cache: Optional[Dict[SampleId, Tuple[List[SampleId], int]]] = field(
-        default=None, 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 = cast(List[SampleId], [i for i in range(self.length)])
+            self.sample_ids = [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._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)}"
+            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)}"
         else:
             raise Exception("length is deprecated, please use sample_ids instead.")
 
@@ -100,36 +79,9 @@ 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
@@ -300,6 +252,7 @@ class MetadataHandler:
 class CustomLatentSpaceHandler:
     function: SectionCallableInterface
     name: str = 'custom_latent_space'
+    use_ls_for_analysis: bool = False
 
 
 # How a chain's latent-space vectors are derived from its per-step forward passes.
@@ -382,7 +335,7 @@ class DatasetIntegrationSetup:
     metrics: List[MetricHandler] = field(default_factory=list)
     instance_metrics: List[InstanceMetricHandler] = field(default_factory=list)
     custom_layers: Dict[str, CustomLayerHandler] = field(default_factory=dict)
-    custom_latent_space: Optional[CustomLatentSpaceHandler] = None
+    custom_latent_spaces: Dict[str, CustomLatentSpaceHandler] = field(default_factory=dict)
     simulations: List[SimulationHandler] = field(default_factory=list)
     autoregressive_step: Optional[AutoregressiveStepHandler] = None
     autoregressive_metrics: List[AutoregressiveMetricHandler] = field(default_factory=list)
@@ -398,6 +351,6 @@ class DatasetSample:
     metadata_is_none: Dict[str, bool]
     index: Union[int, str]
     state: DataStateEnum
-    custom_latent_space: Optional[npt.NDArray[np.float32]] = None
+    custom_latent_spaces: Optional[Dict[str, npt.NDArray[np.float32]]] = None
     instance_masks: Optional[Dict[str, ElementInstance]] = None
 

code_loader/inner_leap_binder/leapbinder.py

@@ -1,9 +1,6 @@
 
-import builtins
 import inspect
-import os
-from contextlib import contextmanager
-from typing import Callable, List, Optional, Dict, Any, Type, Union, get_args, cast, Iterator, Set
+from typing import Callable, List, Optional, Dict, Any, Type, Union, get_args
 
 import numpy as np
 import numpy.typing as npt
@@ -19,8 +16,7 @@ from code_loader.contract.datasetclasses import SectionCallableInterface, InputH
     SimulationHandler, _simulation_context, AutoregressiveStepHandler, AutoregressiveStepCallableInterface, \
     AUTOREGRESSIVE_LATENT_SPACE_AGGREGATIONS, AUTOREGRESSIVE_IMPLICIT_ARG_NAMES, \
     AutoregressiveMetricHandler, AutoregressiveLossHandler, AutoregressiveVisualizerHandler
-from code_loader.contract.enums import LeapDataType, DataStateEnum, DataStateType, MetricDirection, DatasetMetadataType, \
-    TestingSectionEnum
+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
 from code_loader.contract.visualizer_classes import map_leap_data_type_to_visualizer_class
@@ -37,25 +33,6 @@ from code_loader.visualizers.default_visualizers import DefaultVisualizer, \
 mapping_runtime_mode_env_var_mame = '__MAPPING_RUNTIME_MODE__'
 
 
-@contextmanager
-def _track_opened_files() -> Iterator[Set[str]]:
-    opened: Set[str] = set()
-    original_open = builtins.open
-
-    def tracking_open(file: Any, *args: Any, **kwargs: Any) -> Any:
-        mode = args[0] if args else kwargs.get('mode', 'r')
-        is_read_mode = isinstance(mode, str) and not any(c in mode for c in ('a', 'w', 'x'))
-        if is_read_mode and isinstance(file, (str, os.PathLike)):
-            opened.add(os.fspath(file))
-        return original_open(file, *args, **kwargs)
-
-    builtins.open = tracking_open
-    try:
-        yield opened
-    finally:
-        builtins.open = original_open
-
-
 def _stringized_annotation_type_name(annotation: Any) -> Optional[str]:
     """Bare type name if ``annotation`` is a stringized annotation (from ``from __future__
     import annotations`` or a quoted hint), else ``None``. code_loader inspects raw
@@ -536,18 +513,54 @@ class LeapBinder:
         """
         self.setup_container.metadata.append(MetadataHandler(name, function, metadata_type))
 
-    def set_custom_latent_space(self, function: SectionCallableInterface) -> None:
+    def set_custom_latent_space(self, function: SectionCallableInterface,
+                                name: Optional[str] = None,
+                                use_ls_for_analysis: bool = False) -> None:
         """
-        Set a custom latent space function.
+        Register a custom latent space function.
+
+        Multiple custom latent spaces may be registered as long as each has a
+        unique name. They are stored name-keyed (insertion order preserved), so a
+        later registration no longer overrides an earlier one.
 
         Args:
-        function (SectionCallableInterface): The metadata handler function.
+        function (SectionCallableInterface): The latent-space handler function.
             This function receives:
                 subset (PreprocessResponse): The subset of the data.
                 index (int): The index of the sample within the subset.
-            This function should numpy float32 array contains the latent space vec of the sample.
+            This function should return a numpy float32 array containing the latent
+            space vec of the sample.
+        name (Optional[str]): Unique name for this custom latent space. Defaults to
+            the reserved single-LS name for backward compatibility.
+        use_ls_for_analysis (bool): When True, the engine uses this custom latent
+            space for the Out-Of-Distribution and Domain-Gap insights instead of the
+            built-in defaults. At most one registered custom latent space may set this;
+            registering a second one with the flag raises.
         """
-        self.setup_container.custom_latent_space = CustomLatentSpaceHandler(function)
+        if name is None:
+            name = custom_latent_space_attribute
+        if name in self.setup_container.custom_latent_spaces:
+            raise Exception(
+                f"A custom latent space named '{name}' is already registered. Each "
+                f"@tensorleap_custom_latent_space must have a unique name "
+                f"(pass name='...' to distinguish them)."
+            )
+        if use_ls_for_analysis:
+            already_flagged = [
+                existing_name
+                for existing_name, handler in self.setup_container.custom_latent_spaces.items()
+                if handler.use_ls_for_analysis
+            ]
+            if already_flagged:
+                raise Exception(
+                    f"use_ls_for_analysis=True is already set on custom latent space "
+                    f"'{already_flagged[0]}'. Only one custom latent space may set "
+                    f"use_ls_for_analysis=True (it selects the latent space used for the "
+                    f"Out-Of-Distribution and Domain-Gap insights). Set it on '{name}' "
+                    f"or '{already_flagged[0]}', not both."
+                )
+        self.setup_container.custom_latent_spaces[name] = CustomLatentSpaceHandler(
+            function, name, use_ls_for_analysis)
 
     def set_autoregressive_step(self, function: AutoregressiveStepCallableInterface,
                                 latent_space_aggregation: str = 'last_step') -> None:
@@ -572,9 +585,7 @@ class LeapBinder:
         # Builtin chain metadata, declared at parse time so it survives the reporter's
         # metadata type mapping; the placeholder values are overwritten by the engine when a
         # chain finalizes (realized length, truncated-by-safety-cap flag).
-        def builtin_chain_metadata(idx: Any, preprocess: PreprocessResponse) -> Any:
-            if isinstance(idx, list):
-                return [{'length': -1, 'truncated': False} for _ in idx]
+        def builtin_chain_metadata(idx: Any, preprocess: PreprocessResponse) -> Dict[str, Any]:
             return {'length': -1, 'truncated': False}
 
         self.set_metadata(builtin_chain_metadata, 'builtin_chain')
@@ -738,15 +749,6 @@ 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]:
@@ -767,34 +769,7 @@ 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
-        opened_files: Optional[Set[str]] = None
-        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]
-            with _track_opened_files() as opened_files:
-                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)
-
-        def _warn_if_group_spans_multiple_files(payloads: List[DatasetTestResultPayload]) -> None:
-            if opened_files is None or len(opened_files) <= 1:
-                return
-            shown = sorted(opened_files)[:5]
-            suffix = f" (+{len(opened_files) - 5} more)" if len(opened_files) > 5 else ""
-            warning = (
-                f"Encoding declared group 0 (size {len(cast(List[Any], group))}) for '{dataset_base_handler.name}' opened "
-                f"{len(opened_files)} distinct files: {shown}{suffix}. If this group is meant to be one "
-                f"physical source (e.g. one Parquet file), it may be grouped incorrectly - check your "
-                f"preprocess grouping logic. (Only files opened via Python's open() are tracked, so this "
-                f"can miss other I/O paths or be a false positive if grouping isn't meant to reflect file "
-                f"locality.)")
-            for payload in payloads:
-                existing = payload.display.get(TestingSectionEnum.Errors.name, '')
-                payload.display[TestingSectionEnum.Errors.name] = (existing + '\n' if existing else '') + warning
-
+        raw_result = dataset_base_handler.function(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):
@@ -834,7 +809,6 @@ class LeapBinder:
             else:
                 if raw_result is None:
                     if state != DataStateEnum.training:
-                        _warn_if_group_spans_multiple_files(test_result)
                         return test_result
 
                     if dataset_base_handler.metadata_type is None:
@@ -856,7 +830,6 @@ class LeapBinder:
             # setting shape in setup for all encoders
             if isinstance(dataset_base_handler, (InputHandler, GroundTruthHandler)):
                 dataset_base_handler.shape = result_shape
-        _warn_if_group_spans_multiple_files(test_result)
         return test_result
 
     def check_handlers(self, preprocess_result: Dict[DataStateEnum, PreprocessResponse]) -> None:
@@ -892,11 +865,10 @@ 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 = cast(List[Union[int, str]], [0])
-            sim_sample_ids = cast(List[Union[int, str]], preprocess_response.sample_ids)
+            preprocess_response.sample_ids = [0]
             for handler in self.setup_container.inputs:
-                out1 = handler.function(sim_sample_ids[0], preprocess_response)
-                out2 = handler.function(sim_sample_ids[0], preprocess_response)
+                out1 = handler.function(preprocess_response.sample_ids[0], preprocess_response)
+                out2 = handler.function(preprocess_response.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(
@@ -954,7 +926,7 @@ class LeapBinder:
                 "Element instances are not supported together with tensorleap_autoregressive_step: "
                 "instance generation masks the sample's encoded inputs, which do not exist in an "
                 "autoregressive integration. Remove the instance encoders or the autoregressive hook.")
-        if self.setup_container.custom_latent_space is not None:
+        if self.setup_container.custom_latent_spaces:
             raise Exception(
                 "tensorleap_custom_latent_space is not supported together with "
                 "tensorleap_autoregressive_step: the custom latent space function only sees "

code_loader/inner_leap_binder/leapbinder_decorators.py

@@ -40,7 +40,6 @@ _called_from_inside_tl_integration_test_decorator = False
 # to tell a raise in the test body (report as code-flow failure) from a module-level crash
 # before any test ran (report as a plain script crash).
 _integration_test_started = False
-_mapping_dataset_is_grouped = False
 _call_from_tl_platform = os.environ.get('IS_TENSORLEAP_PLATFORM') == 'true'
 
 # ---- warnings store (module-level) ----
@@ -214,53 +213,6 @@ def batch_warning(result, func_name):
         )
 
 
-def _validate_id_or_group(sample_id, preprocess_response, func_name):
-    """Accept either a single sample id (flat dataset) or a group — a list of ids (grouped
-    dataset). The argument shape must match the dataset: a group requires a grouped
-    PreprocessResponse and a single id requires a flat one. Every id must match the declared
-    sample_id_type."""
-    is_group = isinstance(sample_id, list)
-    if is_group and not preprocess_response.is_grouped:
-        raise AssertionError(
-            f'{func_name}() validation failed: got a group (list of sample ids) but the '
-            f'PreprocessResponse is not grouped. Pass a single sample id for a flat dataset.')
-    if not is_group and preprocess_response.is_grouped:
-        raise AssertionError(
-            f'{func_name}() validation failed: got a single sample id but the PreprocessResponse '
-            f'is grouped. Pass a group (list of sample ids) for a grouped dataset.')
-    ids = sample_id if is_group else [sample_id]
-    for sid in ids:
-        assert type(sid) == preprocess_response.sample_id_type, \
-            (f'{func_name}() validation failed: '
-             f'Argument sample_id should be as the same type as defined in the preprocess response '
-             f'{preprocess_response.sample_id_type}. Got {type(sid)}.')
-
-
-def _validate_grouped_result(result, group_size, func_name, validate_single):
-    """A grouped (group) encoder result must be either a single ndarray with leading dim
-    == group_size, or a list of `group_size` per-sample results (each validated by
-    validate_single). Metadata validates its grouped result separately (list only).
-    An `(B, *)` array is validated per-sample slice (not as the whole batched array) so
-    per-sample checks (e.g. channel_dim <= rank) use the sample shape, not the batch shape."""
-    if isinstance(result, np.ndarray):
-        assert len(result.shape) > 0 and result.shape[0] == group_size, \
-            (f'{func_name}() validation failed: expected a grouped output for a group of '
-             f'{group_size}: an array with leading dim {group_size}, got shape {result.shape}.')
-        for single in result:
-            validate_single(single)
-    elif isinstance(result, list):
-        assert len(result) == group_size, \
-            (f'{func_name}() validation failed: expected a grouped output for a group of '
-             f'{group_size}: a list of {group_size} arrays, got {len(result)}.')
-        for r in result:
-            validate_single(r)
-    else:
-        raise AssertionError(
-            f'{func_name}() validation failed: expected a grouped output for a group of '
-            f'{group_size}: either an array with leading dim {group_size} or a list of '
-            f'{group_size} arrays, got {type(result)}.')
-
-
 def _add_mapping_connection(user_unique_name, connection_destinations, arg_names, name, node_mapping_type):
     connection_destinations = [connection_destination for connection_destination in connection_destinations
                                if not isinstance(connection_destination, SamplePreprocessResponse)]
@@ -338,7 +290,11 @@ def tensorleap_integration_test():
 
         def _validate_input_args(*args, **kwargs):
             sample_id, preprocess_response = args
-            _validate_id_or_group(sample_id, preprocess_response, 'tensorleap_integration_test')
+            assert type(sample_id) == preprocess_response.sample_id_type, (
+                f"tensorleap_integration_test validation failed: "
+                f"sample_id type ({type(sample_id).__name__}) does not match the expected "
+                f"type ({preprocess_response.sample_id_type}) from the PreprocessResponse."
+            )
 
         def inner(*args, **kwargs):
             if not _call_from_tl_platform:
@@ -349,7 +305,7 @@ def tensorleap_integration_test():
             expected_names = inspect.getfullargspec(integration_test_function)[0][:2]
             if len(expected_names) < 2:
                 expected_names = ['idx', 'preprocess']
-            validate_args_structure(*args, types_order=[Union[int, str, list], PreprocessResponse],
+            validate_args_structure(*args, types_order=[Union[int, str], PreprocessResponse],
                                     func_name='integration_test', expected_names=expected_names,
                                     **kwargs)
             sample_id, preprocess_response = (
@@ -373,8 +329,6 @@ def tensorleap_integration_test():
                 _reset_model_loop_state()
                 ret = integration_test_function(*args, **kwargs)
 
-                global _mapping_dataset_is_grouped
-                _mapping_dataset_is_grouped = args[1].is_grouped
                 try:
                     os.environ[mapping_runtime_mode_env_var_mame] = 'True'
                     _reset_model_loop_state()
@@ -395,7 +349,6 @@ def tensorleap_integration_test():
                                         f'Integration test is only allowed to call Tensorleap decorators. '
                                         f'Ensure any arithmetics, external library use, Python logic is placed within Tensorleap decoders')
                 finally:
-                    _mapping_dataset_is_grouped = False
                     if mapping_runtime_mode_env_var_mame in os.environ:
                         del os.environ[mapping_runtime_mode_env_var_mame]
             finally:
@@ -1608,10 +1561,13 @@ def tensorleap_metadata(
                 raise Exception(f'Metadata with name {name} already exists. '
                                 f'Please choose another')
 
-        def _validate_input_args(sample_id: Union[int, str, list], preprocess_response: PreprocessResponse):
-            _validate_id_or_group(sample_id, preprocess_response, user_function.__name__)
+        def _validate_input_args(sample_id: Union[int, str], preprocess_response: PreprocessResponse):
+            assert type(sample_id) == preprocess_response.sample_id_type, \
+                (f'{user_function.__name__}() validation failed: '
+                 f'Argument sample_id should be as the same type as defined in the preprocess response '
+                 f'{preprocess_response.sample_id_type}. Got {type(sample_id)}.')
 
-        def _validate_single(result):
+        def _validate_result(result):
             supported_result_types = (type(None), int, str, bool, float, dict, np.floating,
                                       np.bool_, np.unsignedinteger, np.signedinteger, np.integer)
             if isinstance(result, tuple):
@@ -1629,20 +1585,6 @@ def tensorleap_metadata(
                         (f'{user_function.__name__}() validation failed: '
                          f'Values in the return dict should be of type {str(supported_result_types)}. Got {type(value)}.')
 
-        def _validate_result(result, grouped=False, group_size=None):
-            if not grouped:
-                _validate_single(result)
-                return
-            # Grouped metadata: a list of `group_size` per-sample scalars/dicts.
-            assert isinstance(result, list), \
-                (f'{user_function.__name__}() validation failed: expected a grouped output for a group '
-                 f'of {group_size}: a list of {group_size} metadata values, got {type(result)}.')
-            assert len(result) == group_size, \
-                (f'{user_function.__name__}() validation failed: expected a grouped output for a group '
-                 f'of {group_size}: a list of {group_size} metadata values, got {len(result)}.')
-            for value in result:
-                _validate_single(value)
-
         def inner_without_validate(sample_id, preprocess_response):
 
             global _called_from_inside_tl_decorator
@@ -1662,16 +1604,14 @@ def tensorleap_metadata(
                 set_current('tensorleap_metadata')
             if os.environ.get(mapping_runtime_mode_env_var_mame):
                 return None
-            validate_args_structure(*args, types_order=[Union[int, str, list], PreprocessResponse],
+            validate_args_structure(*args, types_order=[Union[int, str], PreprocessResponse],
                                     func_name=user_function.__name__, expected_names=["idx", "preprocess"], **kwargs)
             sample_id, preprocess_response = args if len(args) != 0 else kwargs.values()
             _validate_input_args(sample_id, preprocess_response)
 
-            grouped = isinstance(sample_id, list)
-            group_size = len(sample_id) if grouped else None
             result = inner_without_validate(sample_id, preprocess_response)
 
-            _validate_result(result, grouped, group_size)
+            _validate_result(result)
             if not _call_from_tl_platform:
                 update_env_params_func("tensorleap_metadata", "v")
             return result
@@ -1681,52 +1621,42 @@ def tensorleap_metadata(
     return decorating_function
 
 
-def tensorleap_custom_latent_space():
+def tensorleap_custom_latent_space(name: Optional[str] = None, use_ls_for_analysis: bool = False):
+    assert isinstance(use_ls_for_analysis, bool), \
+        ("tensorleap_custom_latent_space validation failed: use_ls_for_analysis must be a bool. "
+         f"Got {type(use_ls_for_analysis)}.")
+
     def decorating_function(user_function: SectionCallableInterface):
-        def _validate_input_args(sample_id: Union[int, str, list], preprocess_response: PreprocessResponse):
-            _validate_id_or_group(sample_id, preprocess_response, 'tensorleap_custom_latent_space')
+        ls_name = name if name is not None else user_function.__name__
+        def _validate_input_args(sample_id: Union[int, str], preprocess_response: PreprocessResponse):
+            assert isinstance(sample_id, (int, str)), \
+                (f'tensorleap_custom_latent_space validation failed: '
+                 f'Argument sample_id should be either int or str. Got {type(sample_id)}.')
+            assert isinstance(preprocess_response, PreprocessResponse), \
+                (f'tensorleap_custom_latent_space validation failed: '
+                 f'Argument preprocess_response should be a PreprocessResponse. Got {type(preprocess_response)}.')
+            assert type(sample_id) == preprocess_response.sample_id_type, \
+                (f'tensorleap_custom_latent_space validation failed: '
+                 f'Argument sample_id should be as the same type as defined in the preprocess response '
+                 f'{preprocess_response.sample_id_type}. Got {type(sample_id)}.')
 
-        def _validate_single(single_result):
-            assert isinstance(single_result, np.ndarray), \
+        def _validate_result(result):
+            assert isinstance(result, np.ndarray), \
                 (f'tensorleap_custom_latent_space validation failed: '
-                 f'The return type should be a numpy array. Got {type(single_result)}.')
-            if single_result.ndim > 1:
-                flat_dim = int(np.prod(single_result.shape))
+                 f'The return type should be a numpy array. Got {type(result)}.')
+            if result.ndim > 1:
+                flat_dim = int(np.prod(result.shape))
                 store_general_warning(
-                    key=("tensorleap_custom_latent_space_flatten", tuple(single_result.shape)),
+                    key=("tensorleap_custom_latent_space_flatten", ls_name, tuple(result.shape)),
                     message=(
-                        f"tensorleap_custom_latent_space returned per-sample shape {tuple(single_result.shape)} "
-                        f"(ndim={single_result.ndim}). Tensorleap assumes per-sample shape (d, ...) and will "
+                        f"tensorleap_custom_latent_space '{ls_name}' returned per-sample shape {tuple(result.shape)} "
+                        f"(ndim={result.ndim}). Tensorleap assumes per-sample shape (d, ...) and will "
                         f"flatten to ({flat_dim},) before downstream visualization and clustering. "
                         f"If you want a different aggregation (e.g. global average pooling), do it "
                         f"inside your function."
                     ),
                 )
 
-        def _validate_result(result, grouped=False, group_size=None):
-            if not grouped:
-                _validate_single(result)
-                return
-            # Grouped: (B, *) array or list of B per-sample arrays. Validate each per-sample slice so
-            # the per-sample flatten check/warning uses the sample shape, not the (B, *) batch shape.
-            if isinstance(result, np.ndarray):
-                assert result.ndim >= 1 and result.shape[0] == group_size, \
-                    (f'tensorleap_custom_latent_space validation failed: expected a grouped output for a '
-                     f'group of {group_size}: an array with leading dim {group_size}, got shape {result.shape}.')
-                for single in result:
-                    _validate_single(single)
-            elif isinstance(result, list):
-                assert len(result) == group_size, \
-                    (f'tensorleap_custom_latent_space validation failed: expected a grouped output for a '
-                     f'group of {group_size}: a list of {group_size} arrays, got {len(result)}.')
-                for single in result:
-                    _validate_single(single)
-            else:
-                raise AssertionError(
-                    f'tensorleap_custom_latent_space validation failed: expected a grouped output for a '
-                    f'group of {group_size}: either an array with leading dim {group_size} or a list of '
-                    f'{group_size} arrays, got {type(result)}.')
-
         def inner_without_validate(sample_id, preprocess_response):
             global _called_from_inside_tl_decorator
             _called_from_inside_tl_decorator += 1
@@ -1738,7 +1668,8 @@ def tensorleap_custom_latent_space():
 
             return result
 
-        leap_binder.set_custom_latent_space(inner_without_validate)
+        leap_binder.set_custom_latent_space(inner_without_validate, ls_name,
+                                            use_ls_for_analysis=use_ls_for_analysis)
 
         def inner(sample_id, preprocess_response):
             if os.environ.get(mapping_runtime_mode_env_var_mame):
@@ -1746,11 +1677,9 @@ def tensorleap_custom_latent_space():
 
             _validate_input_args(sample_id, preprocess_response)
 
-            grouped = isinstance(sample_id, list)
-            group_size = len(sample_id) if grouped else None
             result = inner_without_validate(sample_id, preprocess_response)
 
-            _validate_result(result, grouped, group_size)
+            _validate_result(result)
             return result
 
         return inner
@@ -2719,15 +2648,6 @@ def tensorleap_preprocess():
             assert len(set(result)) == len(result), \
                 (f'{user_function.__name__}() validation failed: '
                  f'The return list should not contain duplicate PreprocessResponse objects.')
-            # All-or-none: a dataset is either all grouped (nested sample_ids) or all flat.
-            if len({response.is_grouped for response in result}) != 1:
-                shapes = ', '.join(
-                    f"#{i}={'grouped' if response.is_grouped else 'flat'}"
-                    for i, response in enumerate(result))
-                raise AssertionError(
-                    f'{user_function.__name__}() validation failed: '
-                    f'All PreprocessResponses must be either all grouped (nested sample_ids) '
-                    f'or all flat, but got a mix: {shapes}.')
 
         def inner(*args, **kwargs):
             if not _call_from_tl_platform:
@@ -3065,10 +2985,14 @@ def tensorleap_input_encoder(name: str, channel_dim=_UNSET, model_input_index=No
                 f"channel axis in the batched tensor (batch = axis 0), so 0 is invalid. Use 1 "
                 f"for channel-first (NCHW) and -1 for channel-last (NHWC).")
 
-        def _validate_input_args(sample_id: Union[int, str, list], preprocess_response: PreprocessResponse):
-            _validate_id_or_group(sample_id, preprocess_response, user_function.__name__)
+        def _validate_input_args(sample_id: Union[int, str], preprocess_response: PreprocessResponse):
+            assert type(sample_id) == preprocess_response.sample_id_type, \
+                (f'{user_function.__name__}() validation failed: '
+                 f'Argument sample_id should be as the same type as defined in the preprocess response '
+                 f'{preprocess_response.sample_id_type}. Got {type(sample_id)}.')
 
-        def _validate_single(result):
+        def _validate_result(result):
+            validate_output_structure(result, func_name=user_function.__name__, expected_type_name="np.ndarray")
             assert isinstance(result, np.ndarray), \
                 (f'{user_function.__name__}() validation failed: '
                  f'Unsupported return type. Should be a numpy array. Got {type(result)}.')
@@ -3078,13 +3002,6 @@ def tensorleap_input_encoder(name: str, channel_dim=_UNSET, model_input_index=No
             assert channel_dim - 1 <= len(result.shape), (f'{user_function.__name__}() validation failed: '
                                                           f'The channel_dim ({channel_dim}) should be <= to the rank of the resulting input rank ({len(result.shape)}).')
 
-        def _validate_result(result, grouped=False, group_size=None):
-            if not grouped:
-                validate_output_structure(result, func_name=user_function.__name__, expected_type_name="np.ndarray")
-                _validate_single(result)
-                return
-            _validate_grouped_result(result, group_size, user_function.__name__, _validate_single)
-
         def inner_without_validate(sample_id, preprocess_response):
             global _called_from_inside_tl_decorator
             _called_from_inside_tl_decorator += 1
@@ -3101,29 +3018,18 @@ def tensorleap_input_encoder(name: str, channel_dim=_UNSET, model_input_index=No
         def inner(*args, **kwargs):
             if not _call_from_tl_platform:
                 set_current("tensorleap_input_encoder")
-            validate_args_structure(*args, types_order=[Union[int, str, list], PreprocessResponse],
+            validate_args_structure(*args, types_order=[Union[int, str], PreprocessResponse],
                                     func_name=user_function.__name__, expected_names=["idx", "preprocess"], **kwargs)
             sample_id, preprocess_response = args if len(args) != 0 else kwargs.values()
             _validate_input_args(sample_id, preprocess_response)
 
-            grouped = isinstance(sample_id, list)
-            group_size = len(sample_id) if grouped else None
             result = inner_without_validate(sample_id, preprocess_response)
 
-            _validate_result(result, grouped, group_size)
+            _validate_result(result)
 
             if _called_from_inside_tl_decorator == 0 and _called_from_inside_tl_integration_test_decorator:
-                if grouped:
-                    # A grouped result is a list of per-sample arrays (never stacked, so a
-                    # B=1-only model can be fed one sample at a time). A user may also return a
-                    # supported (B, *) array; split it back into per-sample rows first (mirrors
-                    # LeapLoader._to_grouped_list) so both forms match the runtime path. Then add
-                    # a batch dim to each sample, mirroring the flat path's single-sample expand_dims.
-                    rows = result if isinstance(result, list) else [row for row in np.asarray(result)]
-                    result = [np.expand_dims(r, axis=0) for r in rows]
-                else:
-                    batch_warning(result, user_function.__name__)
-                    result = np.expand_dims(result, axis=0)
+                batch_warning(result, user_function.__name__)
+                result = np.expand_dims(result, axis=0)
                 # Emit integration test event once per test
                 try:
                     emit_integration_event_once(AnalyticsEvent.INPUT_ENCODER_INTEGRATION_TEST, {
@@ -3144,15 +3050,8 @@ def tensorleap_input_encoder(name: str, channel_dim=_UNSET, model_input_index=No
         inner.node_mapping = NodeMapping(name, node_mapping_type)
 
         def mapping_inner(*args, **kwargs):
-            if _mapping_dataset_is_grouped:
-                class TempMapping:
-                    def __getitem__(self, key):
-                        # Indexing a grouped input group (input_list[i]) selects a sample from the
-                        # same input source; the model wiring is identical to the un-grouped input.
-                        return self
-            else:
-                class TempMapping:
-                    pass
+            class TempMapping:
+                pass
 
             ret = TempMapping()
             ret.node_mapping = mapping_inner.node_mapping
@@ -3182,10 +3081,15 @@ def tensorleap_gt_encoder(name: str):
                 raise Exception(f'GT with name {name} already exists. '
                                 f'Please choose another')
 
-        def _validate_input_args(sample_id: Union[int, str, list], preprocess_response: PreprocessResponse):
-            _validate_id_or_group(sample_id, preprocess_response, user_function.__name__)
+        def _validate_input_args(sample_id: Union[int, str], preprocess_response: PreprocessResponse):
+            assert type(sample_id) == preprocess_response.sample_id_type, \
+                (f'{user_function.__name__}() validation failed: '
+                 f'Argument sample_id should be as the same type as defined in the preprocess response '
+                 f'{preprocess_response.sample_id_type}. Got {type(sample_id)}.')
 
-        def _validate_single(result):
+        def _validate_result(result):
+            validate_output_structure(result, func_name=user_function.__name__, expected_type_name="np.ndarray",
+                                      gt_flag=True)
             assert isinstance(result, np.ndarray), \
                 (f'{user_function.__name__}() validation failed: '
                  f'Unsupported return type. Should be a numpy array. Got {type(result)}.')
@@ -3193,14 +3097,6 @@ def tensorleap_gt_encoder(name: str):
                 (f'{user_function.__name__}() validation failed: '
                  f'The return type should be a numpy array of type float32. Got {result.dtype}.')
 
-        def _validate_result(result, grouped=False, group_size=None):
-            if not grouped:
-                validate_output_structure(result, func_name=user_function.__name__, expected_type_name="np.ndarray",
-                                          gt_flag=True)
-                _validate_single(result)
-                return
-            _validate_grouped_result(result, group_size, user_function.__name__, _validate_single)
-
         def inner_without_validate(sample_id, preprocess_response):
             global _called_from_inside_tl_decorator
             _called_from_inside_tl_decorator += 1
@@ -3217,29 +3113,18 @@ def tensorleap_gt_encoder(name: str):
         def inner(*args, **kwargs):
             if not _call_from_tl_platform:
                 set_current("tensorleap_gt_encoder")
-            validate_args_structure(*args, types_order=[Union[int, str, list], PreprocessResponse],
+            validate_args_structure(*args, types_order=[Union[int, str], PreprocessResponse],
                                     func_name=user_function.__name__, expected_names=["idx", "preprocess"], **kwargs)
             sample_id, preprocess_response = args
             _validate_input_args(sample_id, preprocess_response)
 
-            grouped = isinstance(sample_id, list)
-            group_size = len(sample_id) if grouped else None
             result = inner_without_validate(sample_id, preprocess_response)
 
-            _validate_result(result, grouped, group_size)
+            _validate_result(result)
 
             if _called_from_inside_tl_decorator == 0 and _called_from_inside_tl_integration_test_decorator:
-                if grouped:
-                    # A grouped result is a list of per-sample arrays (never stacked, so a
-                    # B=1-only model can be fed one sample at a time). A user may also return a
-                    # supported (B, *) array; split it back into per-sample rows first (mirrors
-                    # LeapLoader._to_grouped_list) so both forms match the runtime path. Then add
-                    # a batch dim to each sample, mirroring the flat path's single-sample expand_dims.
-                    rows = result if isinstance(result, list) else [row for row in np.asarray(result)]
-                    result = [np.expand_dims(r, axis=0) for r in rows]
-                else:
-                    batch_warning(result, user_function.__name__)
-                    result = np.expand_dims(result, axis=0)
+                batch_warning(result, user_function.__name__)
+                result = np.expand_dims(result, axis=0)
                 _register_chain_artifact(result, 'ground_truth')
                 # Emit integration test event once per test
                 try:
@@ -3255,15 +3140,8 @@ def tensorleap_gt_encoder(name: str):
         inner.node_mapping = NodeMapping(name, NodeMappingType.GroundTruth)
 
         def mapping_inner(*args, **kwargs):
-            if _mapping_dataset_is_grouped:
-                class TempMapping:
-                    def __getitem__(self, key):
-                        # Indexing a grouped GT group (gt_list[i]) selects a sample from the same
-                        # ground-truth source; wiring is identical to the un-grouped GT node.
-                        return self
-            else:
-                class TempMapping:
-                    pass
+            class TempMapping:
+                pass
 
             ret = TempMapping()
             ret.node_mapping = mapping_inner.node_mapping

code_loader/leaploader.py

@@ -60,36 +60,14 @@ class LeapLoader(LeapLoaderBase):
 
     @lru_cache()
     def exec_script(self) -> None:
-        from code_loader.inner_leap_binder import leapbinder_decorators as _leap_dec
         try:
             os.environ[mapping_runtime_mode_env_var_mame] = 'TRUE'
             self.evaluate_module()
-
-            # A grouped integration test may index a grouped input/gt (image_list[i]) to wire a
-            # single sample of a group; that indexing is only valid when the mapping graph is
-            # built in grouped mode. Detect grouped-ness so the placeholder pass below matches
-            # the dataset. Preprocess is stubbed while mapping mode is on, so run it with mapping
-            # mode momentarily disabled, and cache it so _preprocess_result() doesn't re-run it.
-            is_grouped = False
-            if global_leap_binder.integration_test_func is not None:
-                del os.environ[mapping_runtime_mode_env_var_mame]
-                try:
-                    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())
-                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
-                    PreprocessResponse(state=DataStateType.training, length=0))
-                global_leap_binder.integration_test_func(None, mapping_preprocess)
+                global_leap_binder.integration_test_func(None, PreprocessResponse(state=DataStateType.training, length=0))
         except TypeError as e:
             import traceback
             global_leap_binder.setup_container = DatasetIntegrationSetup()
@@ -102,7 +80,6 @@ class LeapLoader(LeapLoaderBase):
             raise DatasetScriptException(getattr(e, 'message', repr(e))) from e
         finally:
             # ensure that the environment variable is removed after the script execution
-            _leap_dec._mapping_dataset_is_grouped = False
             if mapping_runtime_mode_env_var_mame in os.environ:
                 del os.environ[mapping_runtime_mode_env_var_mame]
 
@@ -232,7 +209,7 @@ class LeapLoader(LeapLoaderBase):
         additional = self._preprocess_result().get(DataStateEnum.additional)
         if additional is None:
             return set()
-        return set(additional.flat_sample_ids)
+        return set(additional.sample_ids)
 
     def _resolve_synthetic(self, sample_id: Union[int, str],
                            state: Optional[DataStateEnum] = None
@@ -273,20 +250,12 @@ 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_list(latent_fn(group_ids, preprocess_state))[pos]
-            else:
-                custom_latent_space = latent_fn(sample_id, preprocess_state)
+        custom_latent_spaces = self._get_custom_latent_spaces(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),
@@ -294,7 +263,7 @@ class LeapLoader(LeapLoaderBase):
                                metadata_is_none=metadata_is_none,
                                index=sample_id,
                                state=state,
-                               custom_latent_space=custom_latent_space,
+                               custom_latent_spaces=custom_latent_spaces,
                                instance_masks=instance_mask)
         return sample
 
@@ -325,11 +294,7 @@ class LeapLoader(LeapLoaderBase):
                     handler.name, handler_result
                 )
 
-        custom_latent_space = None
-        if global_leap_binder.setup_container.custom_latent_space is not None:
-            custom_latent_space = global_leap_binder.setup_container.custom_latent_space.function(
-                original_sample_id, preprocess
-            )
+        custom_latent_spaces = self._get_custom_latent_spaces(original_sample_id, preprocess)
 
         return DatasetSample(
             inputs=inputs,
@@ -338,7 +303,7 @@ class LeapLoader(LeapLoaderBase):
             metadata_is_none=metadata_is_none,
             index=synthetic_index,
             state=DataStateEnum.additional,
-            custom_latent_space=custom_latent_space,
+            custom_latent_spaces=custom_latent_spaces,
             instance_masks=None,
         )
 
@@ -417,7 +382,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.")
 
@@ -513,10 +478,7 @@ class LeapLoader(LeapLoaderBase):
                 if preprocess_response.sample_ids_to_instance_mappings:
                     raise Exception('Element instances are not supported together with '
                                     'tensorleap_autoregressive_step.')
-                # 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]
+                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:
@@ -597,7 +559,7 @@ class LeapLoader(LeapLoaderBase):
         if handler.input_shapes is None:
             preprocess_result = self._preprocess_result()
             first_state_response = next(iter(preprocess_result.values()))
-            first_sample_id = first_state_response.flat_sample_ids[0]
+            first_sample_id = first_state_response.sample_ids[0]
             first_result = handler.function(first_sample_id, None, None, None,
                                             first_state_response)
             if not isinstance(first_result, tuple) or len(first_result) != 2 or \
@@ -936,57 +898,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_list(raw: Any) -> List[npt.NDArray[np.float32]]:
-        """Normalize a grouped encoder result to a list of per-sample arrays.
-
-        Grouped results are NEVER stacked into a (B, *) batch. Grouping is only a storage
-        read-unit (one file load per group); assembling the model batch (B) is the engine's
-        single responsibility, so it stays the only place samples get batched. Keeping
-        per-sample arrays lets the engine batch to any B — including a B=1-only model — without
-        an intermediate stack/unstack. A user encoder that already returned a stacked (B, *)
-        array is split back into its per-sample rows so the contract is a list either way."""
-        if isinstance(raw, list):
-            return raw
-        return [row for row in np.asarray(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 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:
-                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)
-            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]
-
     def _get_dataset_handlers(self, handlers: Iterable[DatasetBaseHandler],
                               state: DataStateEnum, sample_id: Union[int, str]) -> Dict[str, npt.NDArray[np.float32]]:
         result_agg = {}
@@ -998,100 +909,12 @@ 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_list(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 per-sample lists of length B, metadata as per-key lists
-        of length B, index = the group). Grouped results are never stacked here — batching the
-        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_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))
-                  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_list(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]]:
         inputs = self._get_dataset_handlers(global_leap_binder.setup_container.inputs, state, sample_id)
         autoregressive_handler = global_leap_binder.setup_container.autoregressive_step
@@ -1195,18 +1018,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}'
@@ -1225,66 +1042,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())
@@ -1295,10 +1052,42 @@ class LeapLoader(LeapLoaderBase):
 
         return id_type
 
+    @staticmethod
+    def _get_custom_latent_spaces(
+            sample_id: Union[int, str],
+            preprocess: "PreprocessResponse") -> Optional[Dict[str, npt.NDArray[np.float32]]]:
+        handlers = global_leap_binder.setup_container.custom_latent_spaces
+        if not handlers:
+            return None
+        return {name: handler.function(sample_id, preprocess) for name, handler in handlers.items()}
+
     @lru_cache()
     def has_custom_latent_space_decorator(self) -> bool:
         self.exec_script()
-        return global_leap_binder.setup_container.custom_latent_space is not None
+        return len(global_leap_binder.setup_container.custom_latent_spaces) > 0
+
+    @lru_cache()
+    def get_custom_latent_space_names(self) -> Tuple[str, ...]:
+        """Ordered names of all registered custom latent spaces.
+
+        Registration order is the canonical index mapping consumed by the engine
+        (name i -> user_custom_i). Returns a tuple so the lru_cache value is hashable.
+        """
+        self.exec_script()
+        return tuple(global_leap_binder.setup_container.custom_latent_spaces.keys())
+
+    @lru_cache()
+    def get_custom_latent_space_for_analysis(self) -> Optional[str]:
+        """Name of the custom latent space flagged with use_ls_for_analysis=True, if any.
+
+        The engine uses this latent space for the Out-Of-Distribution and Domain-Gap
+        insights instead of the built-in defaults. Returns None when no custom latent
+        space set the flag. At most one is flagged (enforced at registration time)."""
+        self.exec_script()
+        for name, handler in global_leap_binder.setup_container.custom_latent_spaces.items():
+            if handler.use_ls_for_analysis:
+                return name
+        return None
 
     @lru_cache()
     def has_autoregressive_step(self) -> bool:

code_loader/leaploaderbase.py

@@ -218,6 +218,14 @@ class LeapLoaderBase:
         raise NotImplementedError(f'{type(self).__name__} does not implement '
                                   'autoregressive_visualizer_by_name.')
 
+    @abstractmethod
+    def get_custom_latent_space_names(self) -> Tuple[str, ...]:
+        pass
+
+    @abstractmethod
+    def get_custom_latent_space_for_analysis(self) -> Optional[str]:
+        pass
+
     @abstractmethod
     def get_heatmap_visualizer_raw_vis_input_arg_name(self, visualizer_name: str) -> Optional[str]:
         pass

code_loader/utils.py

@@ -17,13 +17,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(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
 

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

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

{code_loader-1.0.196.dev1.dist-info → code_loader-1.0.197.dev1.dist-info}/RECORD

@@ -1,7 +1,7 @@
 LICENSE,sha256=qIwWjdspQeSMTtnFZBC8MuT-95L02FPvzRUdWFxrwJY,1067
 code_loader/__init__.py,sha256=outxRQ0M-zMfV0QGVJmAed5qWfRmyD0TV6-goEGAzBw,406
 code_loader/contract/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-code_loader/contract/datasetclasses.py,sha256=CX4HO3xiuBUiXaeSBEyjAFV7i1NGS7KZ3JnLGxofH28,16164
+code_loader/contract/datasetclasses.py,sha256=LjYe-Ub_WzjXaloFKxwoL1sT1sWOrFSFU6XXxEoYENs,13826
 code_loader/contract/enums.py,sha256=2q-IV_5g9lLE306DIbWA1c0tn5IhDtxsKxyV1x_Lreg,1671
 code_loader/contract/exceptions.py,sha256=jWqu5i7t-0IG0jGRsKF4DjJdrsdpJjIYpUkN1F4RiyQ,51
 code_loader/contract/mapping.py,sha256=sWJhpng-IkOzQnWQdMT5w2ZZ3X1Z_OOzSwCLXIS7oxE,1446
@@ -21,18 +21,18 @@ code_loader/experiment_api/types.py,sha256=MY8xFARHwdVA7p4dxyhD60ShmttgTvb4qdp1o
 code_loader/experiment_api/utils.py,sha256=XZHtxge12TS4H4-8PjV3sKuhp8Ud6ojAiIzTZJEqBqc,3304
 code_loader/experiment_api/workingspace_config_utils.py,sha256=DLzXQCg4dgTV_YgaSbeTVzq-2ja_SQw4zi7LXwKL9cY,990
 code_loader/inner_leap_binder/__init__.py,sha256=koOlJyMNYzGbEsoIbXathSmQ-L38N_pEXH_HvL7beXU,99
-code_loader/inner_leap_binder/leapbinder.py,sha256=yW9EBA4E7YwWoWaiVELWoSP9uWTQjnPl5PfLbmX4cFk,56617
-code_loader/inner_leap_binder/leapbinder_decorators.py,sha256=GyuyffnSmkELj63q_NBAxa6Z0szP6MsaNxDodwCo-pE,183746
-code_loader/leaploader.py,sha256=yWzSc-ynm_ykfup_tk8mwZX3TFq1ZW-fN18fEDXOrfc,83386
-code_loader/leaploaderbase.py,sha256=PvBU_2OQqBfJLDMsbK1s372954eYr4Y3O4r18y5S7uY,10739
+code_loader/inner_leap_binder/leapbinder.py,sha256=MsXT_1DcbCp7dQavHHe2zTEYwGJkIPTO4xaYPMYgsZg,55223
+code_loader/inner_leap_binder/leapbinder_decorators.py,sha256=InZSgWLrt977fIMbw_3WQE0MWhsQhnafhO3uZuMYc0I,176589
+code_loader/leaploader.py,sha256=2_t0IQPQNp9KcSKs4-_Ji8FK4Av8iouyK0eaRYNoUOM,69344
+code_loader/leaploaderbase.py,sha256=Aa8wCcxojf8EBn_14W5I1gKI4OMBdpdEcNAZ_6E2eV4,10940
 code_loader/mixpanel_tracker.py,sha256=rNwRmFifNbdUoqLQvvhhgpKczWpWiEmd8MfyJe27sxw,9131
 code_loader/plot_functions/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 code_loader/plot_functions/plot_functions.py,sha256=2DC-zlVaN13P4VNx5d8csgs80C6SisaeP1-Kq2LW7iM,16075
 code_loader/plot_functions/visualize.py,sha256=gsBAYYkwMh7jIpJeDMPS8G4CW-pxwx6LznoQIvi4vpo,657
-code_loader/utils.py,sha256=Grrp8OLagrVAOYjt83DDmE0CPfjE22DuOQftgqDASyU,7215
+code_loader/utils.py,sha256=gEJCKpDWf6p9_KoEz-0EGiJqxD6QfPgZN6zdKdtSAz8,6628
 code_loader/visualizers/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 code_loader/visualizers/default_visualizers.py,sha256=onRnLE_TXfgLN4o52hQIOOhUcFexGlqJ3xSpQDVLuZM,2604
-code_loader-1.0.196.dev1.dist-info/LICENSE,sha256=qIwWjdspQeSMTtnFZBC8MuT-95L02FPvzRUdWFxrwJY,1067
-code_loader-1.0.196.dev1.dist-info/METADATA,sha256=qjJUAz8M-nUFhOQa1jAFKQXk2Qm-pdh4zns9jWNEglY,1095
-code_loader-1.0.196.dev1.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-code_loader-1.0.196.dev1.dist-info/RECORD,,
+code_loader-1.0.197.dev1.dist-info/LICENSE,sha256=qIwWjdspQeSMTtnFZBC8MuT-95L02FPvzRUdWFxrwJY,1067
+code_loader-1.0.197.dev1.dist-info/METADATA,sha256=USKyIvkEWJ6ak0J1mWQuFQ61R6BP1t3CUbfBlCDAVSo,1095
+code_loader-1.0.197.dev1.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+code_loader-1.0.197.dev1.dist-info/RECORD,,