code-loader 1.0.203__tar.gz → 1.0.204.dev1__tar.gz

{code_loader-1.0.203 → code_loader-1.0.204.dev1}/PKG-INFO

@@ -1,9 +1,9 @@
-Metadata-Version: 2.4
+Metadata-Version: 2.1
 Name: code-loader
-Version: 1.0.203
+Version: 1.0.204.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,7 +20,6 @@ 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.203 → code_loader-1.0.204.dev1}/code_loader/contract/datasetclasses.py

@@ -147,6 +147,8 @@ class ElementInstance:
 SectionCallableInterface = Callable[[Union[int, str], PreprocessResponse], npt.NDArray[np.float32]]
 InstanceCallableInterface = Callable[[Union[int, str], PreprocessResponse, int], Optional[ElementInstance]]
 InstanceLengthCallableInterface = Callable[[Union[int, str], PreprocessResponse], int]
+InstanceSectionCallableInterface = Callable[
+    [Union[int, str], PreprocessResponse, Optional[int]], npt.NDArray[np.float32]]
 
 # (sample_id, prev_inputs, prev_outputs, state, preprocess) -> (next model inputs | None, state).
 # First call per chain receives prev_inputs=None, prev_outputs=None, state=None and returns the
@@ -310,6 +312,7 @@ class CustomLatentSpaceHandler:
     function: SectionCallableInterface
     name: str = 'custom_latent_space'
     use_ls_for_analysis: bool = False
+    instance_aware: bool = False
 
 
 # How a chain's latent-space vectors are derived from its per-step forward passes.
@@ -318,11 +321,6 @@ class CustomLatentSpaceHandler:
 # dominates the model inputs). 'mean': every latent space is the elementwise mean over all steps.
 AUTOREGRESSIVE_LATENT_SPACE_AGGREGATIONS = ('last_step', 'mean')
 
-# Default axis of a KV-cache tensor that indexes token positions: (batch, heads, position,
-# head_dim), the layout HF-style exports produce. Overridable per integration via
-# tensorleap_autoregressive_step(cache_seq_axis=...).
-DEFAULT_CACHE_SEQ_AXIS = 2
-
 # Reserved argument names of autoregressive metrics/losses/visualizers, fed implicitly by the
 # platform from the finished chain (per-chain dicts, no batch axis). Any other argument is wired
 # to a ground-truth encoder through the integration test.
@@ -337,17 +335,6 @@ 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
-    # Which axis of a cache tensor indexes token positions, i.e. where its capacity is read from.
-    # 2 matches the (batch, heads, position, head_dim) layout every HF-style export uses. Declared
-    # rather than inferred because the platform cannot tell positions from heads or head_dim — in
-    # a (1, 3, 64, 64) tensor two of the dims are 64 — and reading the wrong axis silently caps
-    # chains at the number of heads instead of the cache capacity.
-    cache_seq_axis: int = DEFAULT_CACHE_SEQ_AXIS
 
 
 # Per-chain, unbatched callables: called once per finished chain with the final step's tensors —

{code_loader-1.0.203 → code_loader-1.0.204.dev1}/code_loader/inner_leap_binder/leapbinder.py

@@ -15,10 +15,10 @@ from code_loader.contract.datasetclasses import SectionCallableInterface, InputH
     CustomCallableInterfaceMultiArgs, ConfusionMatrixCallableInterfaceMultiArgs, LeapData, \
     CustomMultipleReturnCallableInterfaceMultiArgs, DatasetBaseHandler, custom_latent_space_attribute, \
     RawInputsForHeatmap, VisualizerHandlerData, MetricHandlerData, CustomLossHandlerData, SamplePreprocessResponse, \
-    ElementInstanceMasksHandler, InstanceCallableInterface, CustomLatentSpaceHandler, InstanceMetricHandler, \
+    ElementInstanceMasksHandler, InstanceCallableInterface, InstanceSectionCallableInterface, \
+    CustomLatentSpaceHandler, InstanceMetricHandler, \
     SimulationHandler, _simulation_context, AutoregressiveStepHandler, AutoregressiveStepCallableInterface, \
     AUTOREGRESSIVE_LATENT_SPACE_AGGREGATIONS, AUTOREGRESSIVE_IMPLICIT_ARG_NAMES, \
-    DEFAULT_CACHE_SEQ_AXIS, \
     AutoregressiveMetricHandler, AutoregressiveLossHandler, AutoregressiveVisualizerHandler
 from code_loader.contract.enums import LeapDataType, DataStateEnum, DataStateType, MetricDirection, DatasetMetadataType, \
     TestingSectionEnum
@@ -542,9 +542,10 @@ class LeapBinder:
                             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,
+    def set_custom_latent_space(self, function: Union[SectionCallableInterface, InstanceSectionCallableInterface],
                                 name: Optional[str] = None,
-                                use_ls_for_analysis: bool = False) -> None:
+                                use_ls_for_analysis: bool = False,
+                                instance_aware: bool = False) -> None:
         """
         Register a custom latent space function.
 
@@ -565,6 +566,8 @@ class LeapBinder:
             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.
+        instance_aware (bool): When True, `function` takes a third `instance_id` argument
+            and is called once per element-instance row instead of once per sample.
         """
         if name is None:
             name = custom_latent_space_attribute
@@ -589,59 +592,10 @@ class LeapBinder:
                     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)
-
-    @staticmethod
-    def _validated_cache_seq_axis(cache_seq_axis: int,
-                                  cache_tensors: Optional[Dict[str, str]]) -> int:
-        """The axis must be a non-negative int, and declaring one without a cache is a mistake
-        worth naming — it reads as if it had an effect."""
-        if not isinstance(cache_seq_axis, int) or isinstance(cache_seq_axis, bool) \
-                or cache_seq_axis < 0:
-            raise Exception(f'tensorleap_autoregressive_step: cache_seq_axis must be a '
-                            f'non-negative int axis index, got {cache_seq_axis!r}.')
-        if not cache_tensors and cache_seq_axis != DEFAULT_CACHE_SEQ_AXIS:
-            raise Exception('tensorleap_autoregressive_step: cache_seq_axis was declared without '
-                            'cache_tensors. The axis describes the declared cache tensors, so it '
-                            'has no meaning on its own.')
-        return cache_seq_axis
+            function=function, name=name, use_ls_for_analysis=use_ls_for_analysis, instance_aware=instance_aware)
 
     def set_autoregressive_step(self, function: AutoregressiveStepCallableInterface,
-                                latent_space_aggregation: str = 'last_step',
-                                cache_tensors: Optional[Dict[str, str]] = None,
-                                cache_seq_axis: int = DEFAULT_CACHE_SEQ_AXIS) -> None:
+                                latent_space_aggregation: str = 'last_step') -> 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)
@@ -649,11 +603,6 @@ 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.
-        cache_seq_axis declares which axis of those tensors indexes token positions, i.e. where
-        the platform reads the cache capacity from (default 2, for (batch, heads, position,
-        head_dim)).
         """
         if self.setup_container.autoregressive_step is not None:
             raise Exception('tensorleap_autoregressive_step is already defined. '
@@ -662,11 +611,8 @@ 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)
-        cache_seq_axis = self._validated_cache_seq_axis(cache_seq_axis, cache_tensors)
         self.setup_container.autoregressive_step = AutoregressiveStepHandler(
-            function, latent_space_aggregation=latent_space_aggregation,
-            cache_tensors=cache_tensors, cache_seq_axis=cache_seq_axis)
+            function, latent_space_aggregation=latent_space_aggregation)
 
         # 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

{code_loader-1.0.203 → code_loader-1.0.204.dev1}/code_loader/inner_leap_binder/leapbinder_decorators.py

@@ -24,7 +24,7 @@ from code_loader.contract.datasetclasses import CustomCallableInterfaceMultiArgs
     CustomMultipleReturnCallableInterfaceMultiArgs, ConfusionMatrixCallableInterfaceMultiArgs, CustomCallableInterface, \
     VisualizerCallableInterface, MetadataSectionCallableInterface, PreprocessResponse, SectionCallableInterface, \
     ConfusionMatrixElement, SamplePreprocessResponse, PredictionTypeHandler, InstanceCallableInterface, ElementInstance, \
-    InstanceLengthCallableInterface, AutoregressiveStepCallableInterface, DEFAULT_CACHE_SEQ_AXIS
+    InstanceLengthCallableInterface, InstanceSectionCallableInterface, AutoregressiveStepCallableInterface
 from code_loader.contract.enums import MetricDirection, LeapDataType, DatasetMetadataType, DataStateType
 from code_loader import leap_binder, LeapLoader
 from code_loader.contract.mapping import NodeMapping, NodeMappingType, NodeConnection
@@ -560,22 +560,12 @@ def tensorleap_load_model(prediction_types: Optional[List[PredictionTypeHandler]
                     return result
 
                 def validate_declared_prediction_types(self, ret):
-                    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 (len(self.prediction_types) == len(ret) if isinstance(ret, list) else 1) and len(
+                            self.prediction_types) != 0:
                         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 "
-                            f"prediction types({len(prediction_types)}) != number of model "
-                            f"outputs({n_outputs}){cache_note}")
+                            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})")
 
                 def _convert_onnx_inputs_to_correct_type(
                         self, float_arrays_inputs: Dict[str, np.ndarray]
@@ -733,22 +723,12 @@ 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()
-                    # 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
+                    for i, (input_key, elem) in enumerate(input_dict.items()):
                         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(slot)
-                        slot += 1
+                            elem.node_mapping.type = _safe_get_item(i)
 
                     return ModelOutputPlaceholder()
 
@@ -1807,6 +1787,73 @@ def tensorleap_custom_latent_space(name: Optional[str] = None, use_ls_for_analys
     return decorating_function
 
 
+def tensorleap_instance_custom_latent_space(name: Optional[str] = None, use_ls_for_analysis: bool = False):
+    assert isinstance(use_ls_for_analysis, bool), \
+        ("tensorleap_instance_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: InstanceSectionCallableInterface):
+        ls_name = name if name is not None else user_function.__name__
+
+        argspec = inspect.getfullargspec(user_function)
+        assert len(argspec.args) == 3, \
+            (f'tensorleap_instance_custom_latent_space validation failed: '
+             f'The function should take exactly 3 positional parameters '
+             f'(sample_id, preprocess_response, instance_id). Got {len(argspec.args)}.')
+
+        def _validate_input_args(sample_id: Union[int, str], preprocess_response: PreprocessResponse,
+                                 instance_id: int):
+            _validate_id_or_group(sample_id, preprocess_response, 'tensorleap_instance_custom_latent_space')
+            assert isinstance(instance_id, (int, np.integer)), \
+                (f'tensorleap_instance_custom_latent_space validation failed: '
+                 f'Argument instance_id should be int. Got {type(instance_id)}.')
+
+        def _validate_result(result):
+            assert isinstance(result, np.ndarray), \
+                (f'tensorleap_instance_custom_latent_space validation failed: '
+                 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_instance_custom_latent_space_flatten", ls_name, tuple(result.shape)),
+                    message=(
+                        f"tensorleap_instance_custom_latent_space '{ls_name}' returned per-sample shape "
+                        f"{tuple(result.shape)} (ndim={result.ndim}). Tensorleap assumes per-sample shape "
+                        f"(d, ...) and will flatten to ({flat_dim},) before downstream visualization and "
+                        f"clustering. If you want a different aggregation, do it inside your function."
+                    ),
+                )
+
+        def inner_without_validate(sample_id, preprocess_response, instance_id):
+            global _called_from_inside_tl_decorator
+            _called_from_inside_tl_decorator += 1
+
+            try:
+                result = user_function(sample_id, preprocess_response, instance_id)
+            finally:
+                _called_from_inside_tl_decorator -= 1
+
+            return result
+
+        leap_binder.set_custom_latent_space(inner_without_validate, ls_name,
+                                            use_ls_for_analysis=use_ls_for_analysis, instance_aware=True)
+
+        def inner(sample_id, preprocess_response, instance_id):
+            if os.environ.get(mapping_runtime_mode_env_var_mame):
+                return None
+
+            _validate_input_args(sample_id, preprocess_response, instance_id)
+
+            result = inner_without_validate(sample_id, preprocess_response, instance_id)
+
+            _validate_result(result)
+            return result
+
+        return inner
+
+    return decorating_function
+
+
 _MODEL_LOOP_MAX_STEPS = 1000
 
 _active_model_loop = None
@@ -1904,16 +1951,11 @@ class _MappingStatePlaceholder:
 
 
 class _ModelLoopContext:
-    def __init__(self, sample_id, preprocess_response, prediction_names, is_mapping,
-                 cache_tensors=None):
+    def __init__(self, sample_id, preprocess_response, prediction_names, is_mapping):
         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
@@ -1922,9 +1964,6 @@ 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}')
@@ -2009,19 +2048,12 @@ 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 '
-                       f'platform. Move it into the hook.{extra}')
+                       'platform. Move it into the hook.')
         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 '
@@ -2033,24 +2065,12 @@ class _ModelLoopContext:
             self.last_outputs = raw_outputs
             return raw_outputs
         outputs_list = raw_outputs if isinstance(raw_outputs, list) else [raw_outputs]
-        # 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 '')
+        if len(outputs_list) != len(self.prediction_names):
             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{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):])
+                       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)}
         self.last_outputs = named_outputs
         self.last_outputs_fingerprint = _nest_fingerprint(named_outputs)
         return named_outputs
@@ -2072,20 +2092,8 @@ 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.
 
-        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,
-                                   cache_seq_axis: int = DEFAULT_CACHE_SEQ_AXIS):
+def tensorleap_autoregressive_step(latent_space_aggregation: str = 'last_step'):
     """The feedback hook that drives an autoregressive chain.
 
     Signature of the decorated function:
@@ -2109,29 +2117,6 @@ 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.
-
-    cache_seq_axis declares which axis of those cache tensors indexes token positions — the axis
-    whose length IS the cache capacity, which the platform reads to cap chain length. The default
-    2 matches (batch, heads, position, head_dim), the layout HF-style exports produce. Set it
-    when yours differs, e.g. 1 for (batch, position, heads, head_dim). It cannot be inferred: in
-    a (1, 3, 64, 64) tensor two dims are 64, so a wrong guess would silently cap chains at the
-    number of heads and truncate every chain a few steps in.
     """
     assert isinstance(latent_space_aggregation, str), \
         ('tensorleap_autoregressive_step must be called with parentheses: '
@@ -2279,9 +2264,7 @@ 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,
-                                            cache_tensors=cache_tensors,
-                                            cache_seq_axis=cache_seq_axis)
+                                            latent_space_aggregation=latent_space_aggregation)
 
         def inner(*args, **kwargs):
             if not _call_from_tl_platform:
@@ -2481,9 +2464,8 @@ 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,
-                cache_tensors=leap_binder.setup_container.autoregressive_step.cache_tensors)
+            context = _ModelLoopContext(sample_id, preprocess_response, prediction_names,
+                                        is_mapping)
             _active_model_loop = context
             try:
                 result = user_function(_ModelLoopModelProxy(model, context), sample_id,
@@ -2973,7 +2955,12 @@ def tensorleap_element_instance_preprocess(
                     instances_length = instance_length_encoder(sample_id, preprocess_response)
                     instances_ids = [f'{sample_id}_{instance_id}' for instance_id in range(instances_length)]
                     sample_ids_to_instance_mappings[sample_id] = instances_ids
-                    instance_to_sample_ids_mappings[sample_id] = sample_id
+                    # Only instance ids belong in instance_to_sample_ids_mappings: the engine uses
+                    # membership in it as "is this row an element instance?" (custometl.py,
+                    # samplesgenerator.py). Mapping an original id to itself made that test true for
+                    # every image row, routing it into construct_instance_sample, where looking the
+                    # original up in its own instance list raises
+                    # "Index <id> with sample_id: <id> cannot be found!".
                     for idx, instance_id in enumerate(instances_ids):
                         if not found_instance_metadata:
                             element_instance = instance_mask_encoder(sample_id, preprocess_response, idx)

{code_loader-1.0.203 → code_loader-1.0.204.dev1}/code_loader/leaploader.py

@@ -16,8 +16,7 @@ from code_loader.contract.datasetclasses import DatasetSample, DatasetBaseHandle
     PreprocessResponse, VisualizerHandler, LeapData, \
     PredictionTypeHandler, MetadataHandler, CustomLayerHandler, MetricHandler, VisualizerHandlerData, MetricHandlerData, \
     MetricCallableReturnType, CustomLossHandlerData, CustomLossHandler, RawInputsForHeatmap, SamplePreprocessResponse, \
-    ElementInstance, custom_latent_space_attribute, DatasetIntegrationSetup, InstanceMetricHandler, \
-    _simulation_context, DEFAULT_CACHE_SEQ_AXIS
+    ElementInstance, custom_latent_space_attribute, DatasetIntegrationSetup, InstanceMetricHandler, _simulation_context
 from code_loader.contract.enums import DataStateEnum, TestingSectionEnum, DataStateType, DatasetMetadataType
 from code_loader.contract.exceptions import DatasetScriptException
 from code_loader.contract.responsedataclasses import DatasetIntegParseResult, DatasetTestResultPayload, \
@@ -284,7 +283,7 @@ class LeapLoader(LeapLoaderBase):
 
         metadata, metadata_is_none = self.get_metadata(state, sample_id)
 
-        custom_latent_spaces = self._get_custom_latent_spaces(sample_id, preprocess_result[state])
+        custom_latent_spaces = self._get_custom_latent_spaces(sample_id, preprocess_result[state], instance_id)
         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),
@@ -350,6 +349,9 @@ class LeapLoader(LeapLoaderBase):
                 global_leap_binder.validate_autoregressive_setup()
                 preprocess_test_payload = self._check_preprocess()
                 test_payloads.append(preprocess_test_payload)
+                instance_ls_test_payload = self._check_instance_custom_latent_spaces()
+                if instance_ls_test_payload is not None:
+                    test_payloads.append(instance_ls_test_payload)
                 handlers_test_payloads = self._check_handlers()
                 test_payloads.extend(handlers_test_payloads)
                 simulation_test_payloads = self._check_simulations()
@@ -423,6 +425,37 @@ class LeapLoader(LeapLoaderBase):
             test_result.is_passed = False
         return test_result
 
+    def _check_instance_custom_latent_spaces(self) -> Optional[DatasetTestResultPayload]:
+        instance_aware_names = [
+            name for name, handler in global_leap_binder.setup_container.custom_latent_spaces.items()
+            if handler.instance_aware
+        ]
+        if not instance_aware_names:
+            return None
+
+        test_result = DatasetTestResultPayload('instance_custom_latent_space')
+        errors = []
+        if not global_leap_binder.setup_container.instance_masks:
+            errors.append(
+                f"Instance-aware custom latent space(s) {instance_aware_names} require "
+                f"@tensorleap_instances_masks_encoder to be registered."
+            )
+        preprocess_result = self._preprocess_result()
+        has_instance_mappings = any(
+            preprocess_response.instance_to_sample_ids_mappings
+            for preprocess_response in preprocess_result.values()
+        )
+        if not has_instance_mappings:
+            errors.append(
+                f"Instance-aware custom latent space(s) {instance_aware_names} require the preprocess "
+                f"function to be decorated with @tensorleap_element_instance_preprocess("
+                f"instance_length_encoder, instance_mask_encoder)."
+            )
+        if errors:
+            test_result.is_passed = False
+            test_result.display[TestingSectionEnum.Errors.name] = "\n".join(errors)
+        return test_result
+
     def _check_handlers(self) -> List[DatasetTestResultPayload]:
         preprocess_result = self._preprocess_result()
         result_payloads: List[DatasetTestResultPayload] = []
@@ -1108,6 +1141,9 @@ class LeapLoader(LeapLoaderBase):
         latent_handlers = global_leap_binder.setup_container.custom_latent_spaces
         custom_latent_spaces = None
         if latent_handlers:
+            if any(handler.instance_aware for handler in latent_handlers.values()):
+                raise Exception(
+                    "instance-aware custom latent spaces are not supported with grouped preprocess responses")
             custom_latent_spaces = {
                 name: self._to_grouped_list(handler.function(group_ids, preprocess_state))
                 for name, handler in latent_handlers.items()
@@ -1330,17 +1366,29 @@ class LeapLoader(LeapLoaderBase):
     def _get_custom_latent_spaces(
             self,
             sample_id: Union[int, str],
-            preprocess: "PreprocessResponse") -> Optional[Dict[str, npt.NDArray[np.float32]]]:
+            preprocess: "PreprocessResponse",
+            instance_id: Optional[int] = None) -> 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).
+            if any(handler.instance_aware for handler in handlers.values()):
+                raise Exception(
+                    "instance-aware custom latent spaces are not supported with grouped preprocess responses")
             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()}
+        result: Dict[str, npt.NDArray[np.float32]] = {}
+        for name, handler in handlers.items():
+            if handler.instance_aware:
+                if instance_id is None:
+                    continue
+                result[name] = handler.function(sample_id, preprocess, instance_id)
+            else:
+                result[name] = handler.function(sample_id, preprocess)
+        return result
 
     @lru_cache()
     def has_custom_latent_space_decorator(self) -> bool:
@@ -1349,14 +1397,26 @@ class LeapLoader(LeapLoaderBase):
 
     @lru_cache()
     def get_custom_latent_space_names(self) -> Tuple[str, ...]:
-        """Ordered names of all registered custom latent spaces.
+        """Names of all registered custom latent spaces, in registration order.
 
-        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.
+        The engine resolves each name to its `user_custom_i` index by sorted name, not
+        registration order, so this order is not the index mapping. 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_instance_custom_latent_space_names(self) -> Tuple[str, ...]:
+        """Names of registered custom latent spaces that are instance-aware.
+
+        Names, not indices — the engine resolves each name to its `user_custom_i` index
+        by sorted name.
+        """
+        self.exec_script()
+        return tuple(name for name, handler in global_leap_binder.setup_container.custom_latent_spaces.items()
+                    if handler.instance_aware)
+
     @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.
@@ -1383,34 +1443,6 @@ 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)
-
-    @lru_cache()
-    def get_autoregressive_cache_seq_axis(self) -> int:
-        """Which axis of the declared cache tensors indexes token positions.
-
-        The platform reads the cache capacity off this axis to cap chain length. Defaults to 2
-        ((batch, heads, position, head_dim)) — including when nothing is declared, where it is
-        simply unused.
-        """
-        self.exec_script()
-        handler = global_leap_binder.setup_container.autoregressive_step
-        if handler is None:
-            return DEFAULT_CACHE_SEQ_AXIS
-        return int(handler.cache_seq_axis)
-
     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.203 → code_loader-1.0.204.dev1}/code_loader/leaploaderbase.py

@@ -176,16 +176,6 @@ 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_cache_seq_axis(self) -> int:
-        raise NotImplementedError(f'{type(self).__name__} does not implement '
-                                  'get_autoregressive_cache_seq_axis.')
-
     @abstractmethod
     def get_autoregressive_decorator_names(self) -> Dict[str, List[str]]:
         raise NotImplementedError(f'{type(self).__name__} does not implement '

{code_loader-1.0.203 → code_loader-1.0.204.dev1}/pyproject.toml

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