code-loader 1.0.190__tar.gz → 1.0.191.dev0__tar.gz

{code_loader-1.0.190 → code_loader-1.0.191.dev0}/PKG-INFO

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

{code_loader-1.0.190 → code_loader-1.0.191.dev0}/code_loader/contract/datasetclasses.py

@@ -277,7 +277,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)
 
 
@@ -289,6 +289,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.190 → code_loader-1.0.191.dev0}/code_loader/inner_leap_binder/leapbinder.py

@@ -511,18 +511,34 @@ class LeapBinder:
         """
         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) -> 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.
         """
-        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)."
+            )
+        self.setup_container.custom_latent_spaces[name] = CustomLatentSpaceHandler(function, name)
 
     def set_custom_layer(self, custom_layer: Type[Any], name: str, inspect_layer: bool = False,
                          kernel_index: Optional[int] = None, use_custom_latent_space: bool = False) -> None:

{code_loader-1.0.190 → code_loader-1.0.191.dev0}/code_loader/inner_leap_binder/leapbinder_decorators.py

@@ -1568,8 +1568,9 @@ def tensorleap_metadata(
     return decorating_function
 
 
-def tensorleap_custom_latent_space():
+def tensorleap_custom_latent_space(name: Optional[str] = None):
     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], preprocess_response: PreprocessResponse):
             assert isinstance(sample_id, (int, str)), \
                 (f'tensorleap_custom_latent_space validation failed: '
@@ -1589,9 +1590,9 @@ def tensorleap_custom_latent_space():
             if result.ndim > 1:
                 flat_dim = int(np.prod(result.shape))
                 store_general_warning(
-                    key=("tensorleap_custom_latent_space_flatten", tuple(result.shape)),
+                    key=("tensorleap_custom_latent_space_flatten", ls_name, tuple(result.shape)),
                     message=(
-                        f"tensorleap_custom_latent_space returned per-sample shape {tuple(result.shape)} "
+                        f"tensorleap_custom_latent_space '{ls_name}' returned per-sample shape {tuple(result.shape)} "
                         f"(ndim={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 "
@@ -1610,7 +1611,7 @@ 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)
 
         def inner(sample_id, preprocess_response):
             if os.environ.get(mapping_runtime_mode_env_var_mame):

{code_loader-1.0.190 → code_loader-1.0.191.dev0}/code_loader/leaploader.py

@@ -251,11 +251,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:
-            custom_latent_space = global_leap_binder.setup_container.custom_latent_space.function(sample_id,
-                                                                                                  preprocess_result[
-                                                                                                      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),
@@ -263,7 +259,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
 
@@ -294,11 +290,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,
@@ -307,7 +299,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,
         )
 
@@ -874,10 +866,29 @@ class LeapLoader(LeapLoaderBase):
 
         return id_type
 
+    @staticmethod
+    def _get_custom_latent_spaces(
+            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
+        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())
 
     def get_instances_data(self, state: DataStateEnum) -> Tuple[Dict[str, List[str]], Dict[str, str]]:
         preprocess_result = self._preprocess_result()

{code_loader-1.0.190 → code_loader-1.0.191.dev0}/code_loader/leaploaderbase.py

@@ -149,6 +149,10 @@ class LeapLoaderBase:
     def has_custom_latent_space_decorator(self) -> bool:
         pass
 
+    @abstractmethod
+    def get_custom_latent_space_names(self) -> Tuple[str, ...]:
+        pass
+
     @abstractmethod
     def get_heatmap_visualizer_raw_vis_input_arg_name(self, visualizer_name: str) -> Optional[str]:
         pass

{code_loader-1.0.190 → code_loader-1.0.191.dev0}/pyproject.toml

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