code-loader 1.0.197.dev0__py3-none-any.whl → 1.0.197.dev2__py3-none-any.whl

code_loader/leaploader.py

@@ -41,6 +41,11 @@ def _serialize_sim_bounds(bounds) -> dict:
 
 
 class LeapLoader(LeapLoaderBase):
+    # Opt-in ceiling (env var) on samples materialized per get_samples call — the engine sets it to
+    # turn a mistaken whole-group ask into a clear error instead of a silent worker OOM. Unset => no
+    # ceiling. See get_samples and grouped-fetch-oom-bug.md.
+    _grouped_fetch_max_samples_env = "LEAP_GROUPED_FETCH_MAX_SAMPLES_PER_CALL"
+
     def __init__(self, code_path: str, code_entry_name: str):
         super().__init__(code_path, code_entry_name)
 
@@ -60,14 +65,36 @@ 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()
-                global_leap_binder.integration_test_func(None, PreprocessResponse(state=DataStateType.training, length=0))
+                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)
         except TypeError as e:
             import traceback
             global_leap_binder.setup_container = DatasetIntegrationSetup()
@@ -80,6 +107,7 @@ 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]
 
@@ -209,7 +237,7 @@ class LeapLoader(LeapLoaderBase):
         additional = self._preprocess_result().get(DataStateEnum.additional)
         if additional is None:
             return set()
-        return set(additional.sample_ids)
+        return set(additional.flat_sample_ids)
 
     def _resolve_synthetic(self, sample_id: Union[int, str],
                            state: Optional[DataStateEnum] = None
@@ -250,7 +278,7 @@ class LeapLoader(LeapLoaderBase):
             )
 
         preprocess_result = self._preprocess_result()
-        if state == DataStateEnum.unlabeled and sample_id not in preprocess_result[state].sample_ids:
+        if state == DataStateEnum.unlabeled and sample_id not in preprocess_result[state].flat_sample_ids:
             self._preprocess_result(update_unlabeled_preprocess=True)
 
         metadata, metadata_is_none = self.get_metadata(state, sample_id)
@@ -379,10 +407,10 @@ class LeapLoader(LeapLoaderBase):
         test_result = DatasetTestResultPayload('preprocess')
         try:
             preprocess_result = self._preprocess_result()
-            if self.get_sample_id_type() is str:
+            if issubclass(self.get_sample_id_type(), str):
                 max_allowed_item_size = np.dtype('<U256').itemsize
                 for state, preprocess_response in preprocess_result.items():
-                    sample_ids_array = np.array(preprocess_response.sample_ids)
+                    sample_ids_array = np.array(preprocess_response.flat_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.")
 
@@ -443,6 +471,10 @@ class LeapLoader(LeapLoaderBase):
                 preprocess_response.tl_generated = True
                 if preprocess_response.length < 1:
                     raise ValueError("Simulation returned PreprocessResponse with length < 1")
+                if preprocess_response.is_grouped:
+                    raise ValueError(
+                        "Simulation returned a grouped PreprocessResponse; simulations must "
+                        "return a flat (non-grouped) PreprocessResponse.")
                 preprocess_response.sample_ids = [0]
                 for handler in global_leap_binder.setup_container.inputs:
                     out1 = handler.function(preprocess_response.sample_ids[0], preprocess_response)
@@ -478,7 +510,10 @@ class LeapLoader(LeapLoaderBase):
                 if preprocess_response.sample_ids_to_instance_mappings:
                     raise Exception('Element instances are not supported together with '
                                     'tensorleap_autoregressive_step.')
-                sample_id = preprocess_response.sample_ids[0]
+                # Grouped preprocess is fine: grouping only amortizes the shared base reads
+                # (GT/metadata/custom_latent) — the AR hook always drives a single scalar chain
+                # and is never handed a group.
+                sample_id = preprocess_response.flat_sample_ids[0]
                 first_result = handler.function(sample_id, None, None, None, preprocess_response)
                 second_result = handler.function(sample_id, None, None, None, preprocess_response)
                 if not isinstance(first_result, tuple) or len(first_result) != 2:
@@ -559,7 +594,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.sample_ids[0]
+            first_sample_id = first_state_response.flat_sample_ids[0]
             first_result = handler.function(first_sample_id, None, None, None,
                                             first_state_response)
             if not isinstance(first_result, tuple) or len(first_result) != 2 or \
@@ -898,6 +933,57 @@ 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 = {}
@@ -909,12 +995,127 @@ class LeapLoader(LeapLoaderBase):
             sample_id = original_local_id
         else:
             preprocess_state = preprocess_result[state]
+        if preprocess_state.is_grouped:
+            # Grouped dataset, SINGLE-sample fetch: encode only THIS sample, never the whole group.
+            # Grouping is an I/O read-unit, not an encode-unit — the integration's per-file cache
+            # makes the group's file decode once across consecutive per-sample calls, so a singleton
+            # ([sample_id]) is both a valid single-group subset of the get_samples contract and
+            # correct (samples in a group are independent; see _to_grouped_list). Handing the encoder
+            # the whole group here just to index one sample out materialized B encoded tensors per
+            # sample — a driver of the grouped-fetch worker OOM (see grouped-fetch-oom-bug.md).
+            self._locate_group(preprocess_state, sample_id)  # validates the id belongs to a group
+            for handler in handlers:
+                result_agg[handler.name] = self._to_grouped_list(
+                    handler.function([sample_id], preprocess_state))[0]
+            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.")
+        # Opt-in safety net for the grouped-fetch OOM (grouped-fetch-oom-bug.md). This call
+        # materializes len(group_ids) encoded samples at once; on a large group/file that peak is
+        # what kernel-OOM-kills the worker. A group's file is decoded once by the integration's
+        # per-file cache regardless of how the encode is chunked, so the engine should request
+        # memory-bounded sub-chunks of a group (each a valid single-group subset, e.g.
+        # get_samples(state, group_ids[i:j])) rather than the whole group. When the engine sets the
+        # ceiling below, an over-large ask fails here with a message that names the lever instead of
+        # a silent GENERIC_MEMORY_LIMIT kill. Unset => no ceiling (preserves current whole-group callers).
+        max_per_call = os.environ.get(self._grouped_fetch_max_samples_env)
+        if max_per_call is not None and len(group_ids) > int(max_per_call):
+            raise ValueError(
+                f"get_samples asked to materialize {len(group_ids)} samples in a single call, above "
+                f"the configured ceiling {max_per_call} ({self._grouped_fetch_max_samples_env}). The "
+                f"group's file is decoded once via the integration's per-file cache regardless of "
+                f"chunking, so request memory-bounded sub-chunks of the group "
+                f"(get_samples(state, group_ids[i:j])) instead of the whole group. "
+                f"See grouped-fetch-oom-bug.md.")
+        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])
+            if step_zero_per_sample:
+                for key in self._grouped_dict_keys(step_zero_per_sample, autoregressive_handler.name):
+                    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)  # length-B sequence of scalars/dicts
+            if len(rows) > 0 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_spaces are group-aware like the input/GT encoders: hand each the whole group
+        # in a single call so a file-backed latent fn loads the file once, and normalize to per-sample
+        # lists. 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.
+        latent_handlers = global_leap_binder.setup_container.custom_latent_spaces
+        custom_latent_spaces = None
+        if latent_handlers:
+            custom_latent_spaces = {
+                name: self._to_grouped_list(handler.function(group_ids, preprocess_state))
+                for name, handler in latent_handlers.items()
+            }
+
+        return DatasetSample(inputs=inputs, gt=gt, metadata=metadata, metadata_is_none=metadata_is_none,
+                             index=list(group_ids), state=state, custom_latent_spaces=custom_latent_spaces,
+                             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
@@ -1018,12 +1219,25 @@ class LeapLoader(LeapLoaderBase):
             sample_id = original_local_id
         else:
             preprocess_state = preprocess_result[state]
+        if preprocess_state.is_grouped:
+            self._locate_group(preprocess_state, sample_id)  # validates the id belongs to a group
         for handler in global_leap_binder.setup_container.metadata:
             if requested_metadata_names:
                 if not is_metadata_name_starts_with_handler_name(handler):
                     continue
 
-            handler_result = handler.function(sample_id, preprocess_state)
+            if preprocess_state.is_grouped:
+                # Single-sample metadata fetch: encode only THIS sample, never the whole group —
+                # same memory rationale as _get_dataset_handlers / the get_sample latent branch
+                # (grouped-fetch-oom-bug.md). Metadata values are not guaranteed small scalars: a
+                # metadata encoder may return an np.asarray or read the sample's heavy data to
+                # compute a statistic, so materializing all B rows just to index one out is the same
+                # OOM path this fix closes. The batched get_metadata_multiple_samples correctly keeps
+                # the whole-group encode (amortized across many ids via its group_cache); a single
+                # id does not need it. A singleton is a valid single-group subset of the contract.
+                handler_result = handler.function([sample_id], preprocess_state)[0]
+            else:
+                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}'
@@ -1042,6 +1256,66 @@ 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 len(rows) > 0 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())
@@ -1052,13 +1326,19 @@ class LeapLoader(LeapLoaderBase):
 
         return id_type
 
-    @staticmethod
     def _get_custom_latent_spaces(
+            self,
             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
+        if preprocess.is_grouped:
+            # Single-sample fetch: encode this sample only, not the whole group (same
+            # memory rationale as _get_dataset_handlers; see grouped-fetch-oom-bug.md).
+            self._locate_group(preprocess, sample_id)  # validates group membership
+            return {name: self._to_grouped_list(handler.function([sample_id], preprocess))[0]
+                    for name, handler in handlers.items()}
         return {name: handler.function(sample_id, preprocess) for name, handler in handlers.items()}
 
     @lru_cache()
@@ -1076,6 +1356,19 @@ class LeapLoader(LeapLoaderBase):
         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:
         self.exec_script()

code_loader/leaploaderbase.py

@@ -222,6 +222,10 @@ class LeapLoaderBase:
     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,6 +17,16 @@ 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):
+            # Grouped call (idx is a list of sample ids): never stack the per-sample results 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 below.
+            # Normalize each sample independently (a user-stacked (B, *) array is split back into its
+            # rows) so a ragged group — variable per-sample shapes, e.g. token lists of differing
+            # length — is preserved instead of crashing np.array() on an inhomogeneous shape.
+            rows = result if isinstance(result, list) else np.asarray(result)
+            return [np.asarray(sample) for sample in rows]
         numpy_result: npt.NDArray[np.float32] = np.array(result)
         return numpy_result
 

{code_loader-1.0.197.dev0.dist-info → code_loader-1.0.197.dev2.dist-info}/METADATA

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

{code_loader-1.0.197.dev0.dist-info → code_loader-1.0.197.dev2.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=rKitPlVyI9PZdAZHpH1q-0_s4Vs7oPg8Z659XeuAujM,13788
+code_loader/contract/datasetclasses.py,sha256=LrhOsNirP3MRKUDQcI6u0r-GeSGHppdSUNzVEAkADeU,16816
 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=52nLuWLbOU7ljlABRRuS-e-hKscVlR5zxJTm5FPBdk8,54073
-code_loader/inner_leap_binder/leapbinder_decorators.py,sha256=tW8QmjzkloQU37xryr958aVwyrVOm2FoSFL1LH0S8EU,176272
-code_loader/leaploader.py,sha256=0YZhh2dCouQMCQAakqaPa-nRNrNn65q4rrLTwC8yjE0,68695
-code_loader/leaploaderbase.py,sha256=MMuMv2014d9VM5FQgTsbUBL-9HOBJ4h6L67uF_7jhR4,10837
+code_loader/inner_leap_binder/leapbinder.py,sha256=J-EjPz49kHqneaFhmgzid0MRokjxYb7kX6XkErxwtlc,58909
+code_loader/inner_leap_binder/leapbinder_decorators.py,sha256=WE6gpZP8ckJp8OQ4CgWbz8BuF9KLwAeQYwDk-ua2jsY,186309
+code_loader/leaploader.py,sha256=8q7JnsE1jx-PJNJN0oKak7i47X9dzenbz8lk-ezbjpw,88293
+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=gEJCKpDWf6p9_KoEz-0EGiJqxD6QfPgZN6zdKdtSAz8,6628
+code_loader/utils.py,sha256=Yj06BbS1UIjN9Cv2eQBjDfdDbpQlABAg6KwD3sB1BZ0,7496
 code_loader/visualizers/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 code_loader/visualizers/default_visualizers.py,sha256=onRnLE_TXfgLN4o52hQIOOhUcFexGlqJ3xSpQDVLuZM,2604
-code_loader-1.0.197.dev0.dist-info/LICENSE,sha256=qIwWjdspQeSMTtnFZBC8MuT-95L02FPvzRUdWFxrwJY,1067
-code_loader-1.0.197.dev0.dist-info/METADATA,sha256=WZ4kCsTIDPch7g1xrPJ5fy1EHEQoRGz9q_iygXHTM4s,1095
-code_loader-1.0.197.dev0.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-code_loader-1.0.197.dev0.dist-info/RECORD,,
+code_loader-1.0.197.dev2.dist-info/LICENSE,sha256=qIwWjdspQeSMTtnFZBC8MuT-95L02FPvzRUdWFxrwJY,1067
+code_loader-1.0.197.dev2.dist-info/METADATA,sha256=Ep3bwfV6ttp9_4BD5nL8a5zAzu0C3hauxJIkLisVuDU,1095
+code_loader-1.0.197.dev2.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+code_loader-1.0.197.dev2.dist-info/RECORD,,