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

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

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

{code_loader-1.0.203.dev0 → 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.

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

@@ -15,7 +15,8 @@ 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, \
     AutoregressiveMetricHandler, AutoregressiveLossHandler, AutoregressiveVisualizerHandler
@@ -541,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.
 
@@ -564,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
@@ -588,7 +592,7 @@ class LeapBinder:
                     f"or '{already_flagged[0]}', not both."
                 )
         self.setup_container.custom_latent_spaces[name] = CustomLatentSpaceHandler(
-            function, name, use_ls_for_analysis)
+            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') -> None:
@@ -1017,18 +1021,17 @@ class LeapBinder:
 
         Each entry must be one of:
           - "foreground" (the friendly name for the max-pooled default LS)
-          - "balanced" (the friendly name for the universal fallback default LS),
-            allowed only when a custom latent space is registered with
-            use_ls_for_analysis=True — that flagged LS is what replaces it
           - An input name registered via set_input(name=...)
           - A prediction name registered via add_prediction(name=...)
 
         Forbidden tokens (raise with a clear message):
+          - "balanced" — the friendly name for the universal fallback default
+            LS; cannot be disabled.
           - "user_custom" — populated by set_custom_latent_space; if the user
             registered a custom LS they should use it, not ignore it.
 
-        Order-independent: this runs after all set_input / add_prediction /
-        set_custom_latent_space calls because it consults setup_container.
+        Order-independent: this runs after all set_input / add_prediction calls
+        because it consults setup_container.inputs and prediction_types.
         """
         if not self.leap_analysis_configuration:
             return
@@ -1036,37 +1039,24 @@ class LeapBinder:
         if not names:
             return
 
-        forbidden = {"user_custom"}
+        forbidden = {"balanced", "user_custom"}
         input_names = {h.name for h in self.setup_container.inputs}
         prediction_names = {h.name for h in self.setup_container.prediction_types}
-        allowed_tokens = {"foreground", "balanced"} | input_names | prediction_names
-        analysis_ls_names = [
-            ls_name
-            for ls_name, handler in self.setup_container.custom_latent_spaces.items()
-            if handler.use_ls_for_analysis
-        ]
+        allowed_tokens = {"foreground"} | input_names | prediction_names
 
         for name in names:
             if name in forbidden:
                 raise Exception(
-                    f"Latent space '{name}' cannot be ignored. 'user_custom' is reserved "
-                    f"for set_custom_latent_space."
-                )
-            if name == "balanced" and not analysis_ls_names:
-                raise Exception(
-                    "Latent space 'balanced' (the model-derived default latent space) can "
-                    "only be ignored when a replacement is registered. Add "
-                    "@tensorleap_custom_latent_space(..., use_ls_for_analysis=True) to the "
-                    "latent space that should be analyzed instead, or drop 'balanced' from "
-                    "ignore_latent_spaces — ignoring it with nothing in its place would "
-                    "leave the analyses with no latent space."
+                    f"Latent space '{name}' cannot be ignored. "
+                    f"'balanced' is the universal fallback default LS and 'user_custom' is "
+                    f"reserved for set_custom_latent_space."
                 )
             if name not in allowed_tokens:
                 raise Exception(
                     f"Latent space name '{name}' in ignore_latent_spaces did not match "
                     f"any registered input or prediction. Available names: "
                     f"inputs={sorted(input_names)}, predictions={sorted(prediction_names)}, "
-                    f"tokens=['foreground', 'balanced']."
+                    f"tokens=['foreground']."
                 )
 
     def set_batch_size_to_validate(self, batch_size: int) -> None:

{code_loader-1.0.203.dev0 → 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
+    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
@@ -1787,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
@@ -2888,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.dev0 → code_loader-1.0.204.dev1}/code_loader/leaploader.py

@@ -283,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),
@@ -349,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()
@@ -422,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] = []
@@ -1107,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()
@@ -1329,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:
@@ -1348,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.

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

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