code-loader 1.0.190__py3-none-any.whl → 1.0.191.dev0__py3-none-any.whl

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/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/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/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/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.dist-info → code_loader-1.0.191.dev0.dist-info}/METADATA

@@ -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.dist-info → code_loader-1.0.191.dev0.dist-info}/RECORD

@@ -1,7 +1,7 @@
 LICENSE,sha256=qIwWjdspQeSMTtnFZBC8MuT-95L02FPvzRUdWFxrwJY,1067
 code_loader/__init__.py,sha256=outxRQ0M-zMfV0QGVJmAed5qWfRmyD0TV6-goEGAzBw,406
 code_loader/contract/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-code_loader/contract/datasetclasses.py,sha256=RYZcnX7vLMYhyQtCbX2PscVFNAYULxgAA9U4DRRfLqA,10456
+code_loader/contract/datasetclasses.py,sha256=obtB9Zhzyw241NxerCDn2PkF2MBvcaSZdwHaSMBoM-g,10493
 code_loader/contract/enums.py,sha256=2q-IV_5g9lLE306DIbWA1c0tn5IhDtxsKxyV1x_Lreg,1671
 code_loader/contract/exceptions.py,sha256=jWqu5i7t-0IG0jGRsKF4DjJdrsdpJjIYpUkN1F4RiyQ,51
 code_loader/contract/mapping.py,sha256=sWJhpng-IkOzQnWQdMT5w2ZZ3X1Z_OOzSwCLXIS7oxE,1446
@@ -21,10 +21,10 @@ code_loader/experiment_api/types.py,sha256=MY8xFARHwdVA7p4dxyhD60ShmttgTvb4qdp1o
 code_loader/experiment_api/utils.py,sha256=XZHtxge12TS4H4-8PjV3sKuhp8Ud6ojAiIzTZJEqBqc,3304
 code_loader/experiment_api/workingspace_config_utils.py,sha256=DLzXQCg4dgTV_YgaSbeTVzq-2ja_SQw4zi7LXwKL9cY,990
 code_loader/inner_leap_binder/__init__.py,sha256=koOlJyMNYzGbEsoIbXathSmQ-L38N_pEXH_HvL7beXU,99
-code_loader/inner_leap_binder/leapbinder.py,sha256=PAX9cmT1QQV7PWsbY9rMRWktpM7KryfFlwBiSwSskzg,42654
-code_loader/inner_leap_binder/leapbinder_decorators.py,sha256=k7_kJaEHEuoKWAGRUDXqlZmR1C228xA_Ny9mmliYKys,122098
-code_loader/leaploader.py,sha256=9bUl69X5pk6prgFMCnaoA4MhRoacDb9Ddv3CWgG1F0Y,45968
-code_loader/leaploaderbase.py,sha256=l36qDA00GhZEG5NLKpEtAXgWJA-UQQIhNFGxywK7mUA,6530
+code_loader/inner_leap_binder/leapbinder.py,sha256=t51C-6yH-n8X7VG5YCEPjAreKmfdk1kdseUmspC1KW8,43533
+code_loader/inner_leap_binder/leapbinder_decorators.py,sha256=as8zMdVuKFvWVCDSDi6E1IDj6PU4MbnwuHcYoSUNgn0,122225
+code_loader/leaploader.py,sha256=VGzHpAF0X171Oo3uUlhZ4CbLwu99Ve5brcs3R55XLmI,46293
+code_loader/leaploaderbase.py,sha256=sFMvh7axgoRycLTUWFHsF7Yq20nTLl_k6hgWXe95lg8,6628
 code_loader/mixpanel_tracker.py,sha256=rNwRmFifNbdUoqLQvvhhgpKczWpWiEmd8MfyJe27sxw,9131
 code_loader/plot_functions/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 code_loader/plot_functions/plot_functions.py,sha256=2DC-zlVaN13P4VNx5d8csgs80C6SisaeP1-Kq2LW7iM,16075
@@ -32,7 +32,7 @@ code_loader/plot_functions/visualize.py,sha256=gsBAYYkwMh7jIpJeDMPS8G4CW-pxwx6Lz
 code_loader/utils.py,sha256=YecipkdTA-VcE9F0RQcY9cFnY8P3AksPnHM2Db7xUSk,3972
 code_loader/visualizers/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 code_loader/visualizers/default_visualizers.py,sha256=onRnLE_TXfgLN4o52hQIOOhUcFexGlqJ3xSpQDVLuZM,2604
-code_loader-1.0.190.dist-info/LICENSE,sha256=qIwWjdspQeSMTtnFZBC8MuT-95L02FPvzRUdWFxrwJY,1067
-code_loader-1.0.190.dist-info/METADATA,sha256=Pl7b1dTzYCBphl0mxWCitL0RZ3FnlPfxJhNjEJkcP3k,1090
-code_loader-1.0.190.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-code_loader-1.0.190.dist-info/RECORD,,
+code_loader-1.0.191.dev0.dist-info/LICENSE,sha256=qIwWjdspQeSMTtnFZBC8MuT-95L02FPvzRUdWFxrwJY,1067
+code_loader-1.0.191.dev0.dist-info/METADATA,sha256=OoBB3IpPxf_hKlJTCWdZ3C-HtEH1wZ6Z_5_P_joci5c,1095
+code_loader-1.0.191.dev0.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+code_loader-1.0.191.dev0.dist-info/RECORD,,