code-loader 1.0.202.dev0__tar.gz → 1.0.202.dev2__tar.gz

{code_loader-1.0.202.dev0 → code_loader-1.0.202.dev2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: code-loader
-Version: 1.0.202.dev0
+Version: 1.0.202.dev2
 Summary: 
 License: MIT
 License-File: LICENSE

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

@@ -318,6 +318,11 @@ 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,6 +342,12 @@ class AutoregressiveStepHandler:
     # 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.202.dev0 → code_loader-1.0.202.dev2}/code_loader/inner_leap_binder/leapbinder.py

@@ -18,6 +18,7 @@ from code_loader.contract.datasetclasses import SectionCallableInterface, InputH
     ElementInstanceMasksHandler, InstanceCallableInterface, 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
@@ -622,9 +623,25 @@ class LeapBinder:
             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
+
     def set_autoregressive_step(self, function: AutoregressiveStepCallableInterface,
                                 latent_space_aggregation: str = 'last_step',
-                                cache_tensors: Optional[Dict[str, str]] = None) -> None:
+                                cache_tensors: Optional[Dict[str, str]] = None,
+                                cache_seq_axis: int = DEFAULT_CACHE_SEQ_AXIS) -> 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)
@@ -634,6 +651,9 @@ class LeapBinder:
         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. '
@@ -643,9 +663,10 @@ class LeapBinder:
                             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_tensors=cache_tensors, cache_seq_axis=cache_seq_axis)
 
         # 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.202.dev0 → code_loader-1.0.202.dev2}/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
+    InstanceLengthCallableInterface, AutoregressiveStepCallableInterface, DEFAULT_CACHE_SEQ_AXIS
 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
@@ -733,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()
 
@@ -1893,18 +1903,6 @@ class _MappingStatePlaceholder:
         self._unsupported()
 
 
-def _shares_base(leaf, original):
-    """True when `leaf` is `original` or a view onto it (e.g. original[None])."""
-    if leaf is original:
-        return True
-    base = getattr(leaf, 'base', None)
-    while base is not None:
-        if base is original:
-            return True
-        base = getattr(base, 'base', None)
-    return False
-
-
 class _ModelLoopContext:
     def __init__(self, sample_id, preprocess_response, prediction_names, is_mapping,
                  cache_tensors=None):
@@ -2016,29 +2014,9 @@ class _ModelLoopContext:
         # 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()}
-        # A leaf may be the hook's array itself, or a VIEW of it: `arr[None]` adds the batch
-        # axis an ONNX model needs and shares its memory, so `.base` is the original. Allowing
-        # views keeps the guarantee that matters — no tensor may be COMPUTED between the hook
-        # and the model, since the platform would not see that work — while not forcing every
-        # integration to choose between the identity check and a model that expects a batch
-        # dimension. A genuinely new array has no base among the hook's tensors.
-        def _from_hook(leaf):
-            if id(leaf) in hook_ids:
-                return True
-            base = getattr(leaf, 'base', None)
-            while base is not None:
-                if id(base) in hook_ids:
-                    return True
-                base = getattr(base, 'base', None)
-            return False
-
-        fed_leaves = list(_ndarray_leaves(fed))
-        matched = {id(value) for value in self.fed_inputs.values()
-                   if any(_from_hook(leaf) and (id(leaf) == id(value)
-                                                or _shares_base(leaf, value))
-                          for leaf in fed_leaves)}
-        if not all(_from_hook(leaf) for leaf in fed_leaves) or matched != hook_ids:
+        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 — '
@@ -2106,7 +2084,8 @@ class _ModelLoopModelProxy:
 
 
 def tensorleap_autoregressive_step(latent_space_aggregation: str = 'last_step',
-                                   cache_tensors: Optional[Dict[str, str]] = None):
+                                   cache_tensors: Optional[Dict[str, str]] = None,
+                                   cache_seq_axis: int = DEFAULT_CACHE_SEQ_AXIS):
     """The feedback hook that drives an autoregressive chain.
 
     Signature of the decorated function:
@@ -2146,6 +2125,13 @@ def tensorleap_autoregressive_step(latent_space_aggregation: str = 'last_step',
         })
 
     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: '
@@ -2294,7 +2280,8 @@ def tensorleap_autoregressive_step(latent_space_aggregation: str = 'last_step',
 
         leap_binder.set_autoregressive_step(inner_without_validate,
                                             latent_space_aggregation=latent_space_aggregation,
-                                            cache_tensors=cache_tensors)
+                                            cache_tensors=cache_tensors,
+                                            cache_seq_axis=cache_seq_axis)
 
         def inner(*args, **kwargs):
             if not _call_from_tl_platform:

{code_loader-1.0.202.dev0 → code_loader-1.0.202.dev2}/code_loader/leaploader.py

@@ -16,7 +16,8 @@ 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
+    ElementInstance, custom_latent_space_attribute, DatasetIntegrationSetup, InstanceMetricHandler, \
+    _simulation_context, DEFAULT_CACHE_SEQ_AXIS
 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, \
@@ -1396,6 +1397,20 @@ class LeapLoader(LeapLoaderBase):
             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.202.dev0 → code_loader-1.0.202.dev2}/code_loader/leaploaderbase.py

@@ -181,6 +181,11 @@ class LeapLoaderBase:
         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.202.dev0 → code_loader-1.0.202.dev2}/pyproject.toml

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