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

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

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

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

@@ -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 — '

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

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