code-loader 1.0.201.dev0__tar.gz → 1.0.202.dev1__tar.gz

{code_loader-1.0.201.dev0 → code_loader-1.0.202.dev1}/PKG-INFO

@@ -1,9 +1,9 @@
-Metadata-Version: 2.1
+Metadata-Version: 2.4
 Name: code-loader
-Version: 1.0.201.dev0
+Version: 1.0.202.dev1
 Summary: 
-Home-page: https://github.com/tensorleap/code-loader
 License: MIT
+License-File: LICENSE
 Author: dorhar
 Author-email: doron.harnoy@tensorleap.ai
 Requires-Python: >=3.8,<3.13
@@ -20,6 +20,7 @@ Requires-Dist: numpy (>=2.3.2,<3.0.0) ; python_version >= "3.11" and python_vers
 Requires-Dist: psutil (>=5.9.5,<6.0.0)
 Requires-Dist: pyyaml (>=6.0.2,<7.0.0)
 Requires-Dist: requests (>=2.32.3,<3.0.0)
+Project-URL: Homepage, https://github.com/tensorleap/code-loader
 Project-URL: Repository, https://github.com/tensorleap/code-loader
 Description-Content-Type: text/markdown
 

{code_loader-1.0.201.dev0 → code_loader-1.0.202.dev1}/code_loader/contract/datasetclasses.py

@@ -309,6 +309,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.
@@ -331,6 +332,11 @@ class AutoregressiveStepHandler:
     # time. Fills the role InputHandler.shape plays for input encoders.
     input_shapes: Optional[Dict[str, List[int]]] = None
     latent_space_aggregation: str = 'last_step'
+    # {model_input_name: model_output_name} for a KV-cache (or any threaded state): inputs whose
+    # value is the named output from the previous step. The platform threads these itself — they
+    # are never passed to the hook and never reported — so the hook must NOT return them. Empty
+    # when the model has no cache, in which case the chain re-feeds its full prompt each step.
+    cache_tensors: Optional[Dict[str, str]] = None
 
 
 # Per-chain, unbatched callables: called once per finished chain with the final step's tensors —
@@ -391,7 +397,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)
@@ -407,6 +413,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-1.0.201.dev0 → code_loader-1.0.202.dev1}/code_loader/contract/visualizer_classes.py

@@ -353,7 +353,7 @@ class LeapAudio:
     waveform (a LeapGraph). Build one with `default_audio_visualizer`.
 
     Attributes:
-    audio (npt.NDArray): Mono waveform samples, shaped [N]. Multi-channel input
+    audio (npt.NDArray[np.float32]): Mono waveform samples, shaped [N]. Multi-channel input
         should be mixed down before constructing LeapAudio.
     sample_rate (int): Samples per second — needed to encode the clip and to map
         the playhead position to time.
@@ -362,7 +362,7 @@ class LeapAudio:
         driven by the audio element's currentTime, so no time range is needed here.
     type (LeapDataType): The data type, default is LeapDataType.Audio.
     """
-    audio: npt.NDArray
+    audio: npt.NDArray[np.float32]
     sample_rate: int
     visual: Union[LeapImage, LeapGraph]
     type: LeapDataType = LeapDataType.Audio

{code_loader-1.0.201.dev0 → code_loader-1.0.202.dev1}/code_loader/inner_leap_binder/leapbinder.py

@@ -534,23 +534,97 @@ class LeapBinder:
             leap_binder.set_metadata(metadata_handler_index, name='metadata_index')
             leap_binder.set_metadata(metadata_handler_image_mean, name='metadata_image_mean')
         """
+        if not isinstance(name, str):
+            raise TypeError(f"set_metadata validation failed.\n`name` must be a string, got type "
+                            f"{type(name).__name__}. When using @tensorleap_metadata, the metadata "
+                            f"types go in the second argument: "
+                            f"@tensorleap_metadata('my_metadata', {{'key': DatasetMetadataType.int}}).")
         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)
+
+    @staticmethod
+    def _validated_cache_tensors(cache_tensors: Optional[Dict[str, str]]
+                                 ) -> Optional[Dict[str, str]]:
+        """Structural checks only — shape/dtype/existence are validated by the platform, which
+        is the side that has the model graph. This catches the typos that are cheap to catch
+        here, at integration-test time, instead of at push."""
+        if not cache_tensors:
+            return None
+        if not isinstance(cache_tensors, dict):
+            raise Exception(f'tensorleap_autoregressive_step: cache_tensors must be a dict of '
+                            f'{{model_input_name: model_output_name}}, got '
+                            f'{type(cache_tensors).__name__}.')
+        seen_outputs: Dict[str, str] = {}
+        for input_name, output_name in cache_tensors.items():
+            for role, name in (('key', input_name), ('value', output_name)):
+                if not isinstance(name, str) or not name:
+                    raise Exception(
+                        f'tensorleap_autoregressive_step: cache_tensors {role}s must be '
+                        f'non-empty model tensor names, got {name!r}.')
+            if input_name == output_name:
+                raise Exception(
+                    f'tensorleap_autoregressive_step: cache_tensors maps "{input_name}" to '
+                    f'itself. A model input and output cannot share a name — the mapping is '
+                    f'{{input_name: output_name}}.')
+            if output_name in seen_outputs:
+                raise Exception(
+                    f'tensorleap_autoregressive_step: model output "{output_name}" is declared '
+                    f'as the cache source for both "{seen_outputs[output_name]}" and '
+                    f'"{input_name}". Each cache output feeds exactly one input.')
+            seen_outputs[output_name] = input_name
+        return dict(cache_tensors)
 
     def set_autoregressive_step(self, function: AutoregressiveStepCallableInterface,
-                                latent_space_aggregation: str = 'last_step') -> None:
+                                latent_space_aggregation: str = 'last_step',
+                                cache_tensors: Optional[Dict[str, str]] = None) -> None:
         """
         Set the autoregressive step hook — the feedback function that drives a chain:
         it supplies the model's initial inputs on its first call (prev_inputs=None, prev_outputs=None)
@@ -558,6 +632,8 @@ class LeapBinder:
         An autoregressive integration has no input encoders; the hook is the sole input source.
         latent_space_aggregation declares how the chain's latent-space vectors are derived from
         its steps (see AUTOREGRESSIVE_LATENT_SPACE_AGGREGATIONS).
+        cache_tensors declares a KV-cache as {model_input_name: model_output_name}; the platform
+        threads those tensors itself and the hook neither receives nor returns them.
         """
         if self.setup_container.autoregressive_step is not None:
             raise Exception('tensorleap_autoregressive_step is already defined. '
@@ -566,8 +642,10 @@ class LeapBinder:
             raise Exception(f'tensorleap_autoregressive_step: unknown latent_space_aggregation '
                             f'{latent_space_aggregation!r}. Supported values: '
                             f'{", ".join(AUTOREGRESSIVE_LATENT_SPACE_AGGREGATIONS)}.')
+        cache_tensors = self._validated_cache_tensors(cache_tensors)
         self.setup_container.autoregressive_step = AutoregressiveStepHandler(
-            function, latent_space_aggregation=latent_space_aggregation)
+            function, latent_space_aggregation=latent_space_aggregation,
+            cache_tensors=cache_tensors)
 
         # 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
@@ -958,7 +1036,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-1.0.201.dev0 → code_loader-1.0.202.dev1}/code_loader/inner_leap_binder/leapbinder_decorators.py

@@ -560,12 +560,22 @@ def tensorleap_load_model(prediction_types: Optional[List[PredictionTypeHandler]
                     return result
 
                 def validate_declared_prediction_types(self, ret):
-                    if not (len(self.prediction_types) == len(ret) if isinstance(ret, list) else 1) and len(
-                            self.prediction_types) != 0:
+                    n_outputs = len(ret) if isinstance(ret, list) else 1
+                    # A declared KV-cache is plumbing: those outputs are fed straight back into
+                    # the model and never reported, so they carry no prediction type. Count only
+                    # the rest, or every cached model would demand 60 meaningless declarations.
+                    step = leap_binder.setup_container.autoregressive_step
+                    n_cache = len((step.cache_tensors if step else None) or {})
+                    if len(self.prediction_types) != 0 and \
+                            n_outputs - n_cache != len(self.prediction_types):
                         if not _call_from_tl_platform:
                             update_env_params_func("tensorleap_load_model", "x")
+                        cache_note = (f' ({n_cache} of them are the declared KV-cache, which '
+                                      f'needs none)' if n_cache else '')
                         raise Exception(
-                            f"tensorleap_load_model validation failed: number of declared prediction types({len(prediction_types)}) != number of model outputs({len(ret) if isinstance(ret, list) else 1})")
+                            f"tensorleap_load_model validation failed: number of declared "
+                            f"prediction types({len(prediction_types)}) != number of model "
+                            f"outputs({n_outputs}){cache_note}")
 
                 def _convert_onnx_inputs_to_correct_type(
                         self, float_arrays_inputs: Dict[str, np.ndarray]
@@ -723,12 +733,22 @@ def tensorleap_load_model(prediction_types: Optional[List[PredictionTypeHandler]
                     assert isinstance(input_dict, dict), \
                         f'Expected input_dict to be a dict, got {type(input_dict)} instead.'
                     seen: set = set()
-                    for i, (input_key, elem) in enumerate(input_dict.items()):
+                    # A declared KV-cache is threaded by the platform, not wired from an input
+                    # encoder, so those slots have no mapping node to assign — the loop feeds
+                    # them real arrays even in mapping mode. Skip them, and count only the wired
+                    # inputs so the remaining slot indices stay contiguous.
+                    step_handler = leap_binder.setup_container.autoregressive_step
+                    cache_tensors = (step_handler.cache_tensors if step_handler else None) or {}
+                    slot = 0
+                    for input_key, elem in input_dict.items():
+                        if input_key in cache_tensors:
+                            continue
                         self._reject_reused_input_source(elem, seen)
                         if isinstance(input_key, NodeMappingType):
                             elem.node_mapping.type = input_key
                         else:
-                            elem.node_mapping.type = _safe_get_item(i)
+                            elem.node_mapping.type = _safe_get_item(slot)
+                        slot += 1
 
                     return ModelOutputPlaceholder()
 
@@ -1703,8 +1723,14 @@ 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):
+        ls_name = name if name is not None else user_function.__name__
+
         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')
 
@@ -1715,9 +1741,9 @@ def tensorleap_custom_latent_space():
             if single_result.ndim > 1:
                 flat_dim = int(np.prod(single_result.shape))
                 store_general_warning(
-                    key=("tensorleap_custom_latent_space_flatten", tuple(single_result.shape)),
+                    key=("tensorleap_custom_latent_space_flatten", ls_name, tuple(single_result.shape)),
                     message=(
-                        f"tensorleap_custom_latent_space returned per-sample shape {tuple(single_result.shape)} "
+                        f"tensorleap_custom_latent_space '{ls_name}' returned per-sample shape {tuple(single_result.shape)} "
                         f"(ndim={single_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 "
@@ -1760,7 +1786,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):
@@ -1877,11 +1904,16 @@ class _MappingStatePlaceholder:
 
 
 class _ModelLoopContext:
-    def __init__(self, sample_id, preprocess_response, prediction_names, is_mapping):
+    def __init__(self, sample_id, preprocess_response, prediction_names, is_mapping,
+                 cache_tensors=None):
         self.sample_id = sample_id
         self.preprocess_response = preprocess_response
         self.prediction_names = prediction_names
         self.is_mapping = is_mapping
+        # Declared KV-cache, {model_input_name: model_output_name}. These tensors are supplied
+        # by the loop (and by the platform at runtime), NOT by the hook, so both the fed-inputs
+        # check and the output-count check have to account for them.
+        self.cache_tensors = dict(cache_tensors or {})
         self.phase = 'awaiting_first_hook'
         self.steps = 0
         self.fed_inputs = None
@@ -1890,6 +1922,9 @@ class _ModelLoopContext:
         self.last_outputs_fingerprint = None
         self.last_state = None
         self.last_state_fingerprint = None
+        # Declared-cache outputs of the last model call, in graph order (predictions first, so
+        # this is the tail). Read by the loop through the model proxy.
+        self.last_cache_outputs = []
 
     def _fail(self, message):
         raise LeapValidationError(f'tensorleap_model_loop validation failed: {message}')
@@ -1974,12 +2009,19 @@ class _ModelLoopContext:
                        'next_inputs the loop must exit without further hook or model calls.')
         if self.is_mapping:
             return
+        # A declared KV-cache is threaded by the platform, so the hook neither receives nor
+        # returns it and the loop is expected to supply it. Exclude those keys from the
+        # identity check rather than demanding the hook produce tensors it cannot see.
+        if isinstance(fed, dict) and self.cache_tensors:
+            fed = {key: value for key, value in fed.items() if key not in self.cache_tensors}
         fed_ids = {id(leaf) for leaf in _ndarray_leaves(fed)}
         hook_ids = {id(value) for value in self.fed_inputs.values()}
         if fed_ids != hook_ids:
+            extra = (' (the declared KV-cache inputs are exempt — those are threaded for you)'
+                     if self.cache_tensors else '')
             self._fail('the model must be fed exactly the tensors the last hook call returned — '
                        'any computation between the hook and the model is invisible to the '
-                       'platform. Move it into the hook.')
+                       f'platform. Move it into the hook.{extra}')
         if _nest_fingerprint(self.fed_inputs) != self.fed_inputs_fingerprint:
             self._fail('the model inputs were mutated in place after the hook returned them — '
                        'the platform feeds the model the tensors exactly as the hook returned '
@@ -1991,12 +2033,24 @@ class _ModelLoopContext:
             self.last_outputs = raw_outputs
             return raw_outputs
         outputs_list = raw_outputs if isinstance(raw_outputs, list) else [raw_outputs]
-        if len(outputs_list) != len(self.prediction_names):
+        # Cache outputs are plumbing, not predictions: they are fed straight back into the model
+        # and never reported, so they need no prediction type. Only the remainder is keyed.
+        n_cache_outputs = len(self.cache_tensors)
+        if len(outputs_list) - n_cache_outputs != len(self.prediction_names):
+            cache_note = (f' ({n_cache_outputs} of them are the declared KV-cache, which needs '
+                          f'no prediction types)' if n_cache_outputs else '')
             self._fail(f'the model returned {len(outputs_list)} outputs but '
                        f'{len(self.prediction_names)} prediction types are declared on '
-                       f'tensorleap_load_model — declare one prediction type per model output.')
-        named_outputs = {name: np.asarray(output)
-                         for name, output in zip(self.prediction_names, outputs_list)}
+                       f'tensorleap_load_model{cache_note} — declare one prediction type per '
+                       f'non-cache model output.')
+        # Predictions come first in graph order, as the export declares them; the cache follows.
+        named_outputs = {name: np.asarray(output) for name, output
+                         in zip(self.prediction_names, outputs_list)}
+        # The hook must receive predictions only, exactly as on the platform, and must receive
+        # this very dict (on_hook_call checks identity). The loop still needs the cache, because
+        # locally it stands in for the platform and threads it by hand — so that goes out on the
+        # proxy rather than being mixed into the hook's view.
+        self.last_cache_outputs = list(outputs_list[len(self.prediction_names):])
         self.last_outputs = named_outputs
         self.last_outputs_fingerprint = _nest_fingerprint(named_outputs)
         return named_outputs
@@ -2018,8 +2072,19 @@ class _ModelLoopModelProxy:
     def get_inputs(self):
         return self._model.get_inputs()
 
+    @property
+    def cache_outputs(self):
+        """Declared-cache outputs of the last model call, in graph order.
 
-def tensorleap_autoregressive_step(latent_space_aggregation: str = 'last_step'):
+        Only the local loop needs these: it stands in for the platform, which threads the cache
+        for you at runtime. They are kept off the hook's `prev_outputs` deliberately — the hook
+        must not see the cache in either place, or the local run would diverge from the real one.
+        """
+        return self._context.last_cache_outputs
+
+
+def tensorleap_autoregressive_step(latent_space_aggregation: str = 'last_step',
+                                   cache_tensors: Optional[Dict[str, str]] = None):
     """The feedback hook that drives an autoregressive chain.
 
     Signature of the decorated function:
@@ -2043,6 +2108,22 @@ def tensorleap_autoregressive_step(latent_space_aggregation: str = 'last_step'):
     space comes from the final step's forward pass, except input-kind latent spaces which come
     from the first step — the original sample, before generated content dominates the model
     inputs. 'mean': every latent space is the elementwise mean over all steps of the chain.
+
+    cache_tensors declares a KV-cache (or any threaded state) as
+    {model_input_name: model_output_name} — model inputs whose value is the named output from
+    the previous step. The platform threads those tensors itself and holds them on the GPU, so
+    the hook neither receives them in prev_inputs/prev_outputs nor returns them in next_inputs;
+    return only the inputs you actually compute (e.g. input_ids, position_ids, attention_mask).
+    The cache must be fixed-capacity: the declared input and output must have the same, fully
+    concrete shape, since chains at different steps are batched together and the cache is
+    preallocated. Exports with a growing cache (past length P in, P+1 out) are rejected at push.
+
+        @tensorleap_autoregressive_step(cache_tensors={
+            f'past_key_values.{i}.{kv}': f'present.{i}.{kv}'
+            for i in range(30) for kv in ('key', 'value')
+        })
+
+    Omit it for a model with no cache — the chain then re-feeds its whole prompt every step.
     """
     assert isinstance(latent_space_aggregation, str), \
         ('tensorleap_autoregressive_step must be called with parentheses: '
@@ -2190,7 +2271,8 @@ def tensorleap_autoregressive_step(latent_space_aggregation: str = 'last_step'):
             return result
 
         leap_binder.set_autoregressive_step(inner_without_validate,
-                                            latent_space_aggregation=latent_space_aggregation)
+                                            latent_space_aggregation=latent_space_aggregation,
+                                            cache_tensors=cache_tensors)
 
         def inner(*args, **kwargs):
             if not _call_from_tl_platform:
@@ -2390,8 +2472,9 @@ def tensorleap_model_loop():
             if not prediction_names:
                 _fail('tensorleap_model_loop requires prediction types declared on '
                       'tensorleap_load_model — the loop keys the model outputs by their names.')
-            context = _ModelLoopContext(sample_id, preprocess_response, prediction_names,
-                                        is_mapping)
+            context = _ModelLoopContext(
+                sample_id, preprocess_response, prediction_names, is_mapping,
+                cache_tensors=leap_binder.setup_container.autoregressive_step.cache_tensors)
             _active_model_loop = context
             try:
                 result = user_function(_ModelLoopModelProxy(model, context), sample_id,

{code_loader-1.0.201.dev0 → code_loader-1.0.202.dev1}/code_loader/leaploader.py

@@ -283,17 +283,7 @@ class LeapLoader(LeapLoaderBase):
 
         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:
-                # 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_state, sample_id)  # validates group membership
-                custom_latent_space = self._to_grouped_list(latent_fn([sample_id], preprocess_state))[0]
-            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),
@@ -301,7 +291,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
 
@@ -332,11 +322,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,
@@ -345,7 +331,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,
         )
 
@@ -1114,17 +1100,20 @@ class LeapLoader(LeapLoaderBase):
                 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))
+        # 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_space=custom_latent_space,
+                             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]]:
@@ -1337,10 +1326,48 @@ class LeapLoader(LeapLoaderBase):
 
         return id_type
 
+    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()
     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:
@@ -1355,6 +1382,20 @@ class LeapLoader(LeapLoaderBase):
             return 'last_step'
         return handler.latent_space_aggregation
 
+    @lru_cache()
+    def get_autoregressive_cache_tensors(self) -> Dict[str, str]:
+        """Declared KV-cache pairing, {model_input_name: model_output_name}.
+
+        Empty when nothing is declared — the platform then runs the chain uncached. The platform
+        validates these names against the real model graph at push time; only structural checks
+        happen here.
+        """
+        self.exec_script()
+        handler = global_leap_binder.setup_container.autoregressive_step
+        if handler is None or not handler.cache_tensors:
+            return {}
+        return dict(handler.cache_tensors)
+
     def run_autoregressive_step(self, sample_id: Union[int, str],
                                 prev_inputs: Optional[Dict[str, npt.NDArray[np.float32]]],
                                 prev_outputs: Optional[Dict[str, npt.NDArray[np.float32]]],

{code_loader-1.0.201.dev0 → code_loader-1.0.202.dev1}/code_loader/leaploaderbase.py

@@ -176,6 +176,11 @@ class LeapLoaderBase:
         raise NotImplementedError(f'{type(self).__name__} does not implement '
                                   'get_autoregressive_latent_space_aggregation.')
 
+    @abstractmethod
+    def get_autoregressive_cache_tensors(self) -> Dict[str, str]:
+        raise NotImplementedError(f'{type(self).__name__} does not implement '
+                                  'get_autoregressive_cache_tensors.')
+
     @abstractmethod
     def get_autoregressive_decorator_names(self) -> Dict[str, List[str]]:
         raise NotImplementedError(f'{type(self).__name__} does not implement '
@@ -218,6 +223,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-1.0.201.dev0 → code_loader-1.0.202.dev1}/code_loader/visualizers/default_visualizers.py

@@ -33,7 +33,7 @@ def default_video_visualizer(data: npt.NDArray[np.float32]) -> LeapVideo:
 _MAX_WAVEFORM_POINTS = 2000
 
 
-def _to_mono(audio: npt.NDArray) -> npt.NDArray:
+def _to_mono(audio: npt.NDArray[np.float32]) -> npt.NDArray[np.float32]:
     a = np.squeeze(np.asarray(audio))
     if a.ndim == 2:  # mix down: channels are the shorter axis
         a = a.mean(axis=int(np.argmin(a.shape)))
@@ -42,7 +42,7 @@ def _to_mono(audio: npt.NDArray) -> npt.NDArray:
     return a.astype(np.float32)
 
 
-def _waveform_graph(series: npt.NDArray, duration: float) -> LeapGraph:
+def _waveform_graph(series: npt.NDArray[np.float32], duration: float) -> LeapGraph:
     s = np.squeeze(np.asarray(series)).astype(np.float32)
     # ponytail: recharts chokes on 100k+ points; stride down for the default view. x_range keeps full time span.
     step = max(1, s.shape[0] // _MAX_WAVEFORM_POINTS)
@@ -50,9 +50,9 @@ def _waveform_graph(series: npt.NDArray, duration: float) -> LeapGraph:
     return LeapGraph(s, x_label='Time [s]', y_label='Amplitude', x_range=(0.0, duration))
 
 
-def default_audio_visualizer(audio: npt.NDArray,
+def default_audio_visualizer(audio: npt.NDArray[np.float32],
                              sample_rate: int,
-                             visual: Optional[npt.NDArray] = None,
+                             visual: Optional[npt.NDArray[np.float32]] = None,
                              time_axis_index: int = -1) -> LeapAudio:
     """Package a playable audio sample plus the visual a playhead runs through.
 

{code_loader-1.0.201.dev0 → code_loader-1.0.202.dev1}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "code-loader"
-version = "1.0.201.dev0"
+version = "1.0.202.dev1"
 description = ""
 authors = ["dorhar <doron.harnoy@tensorleap.ai>"]
 license = "MIT"