code-loader 1.0.207.dev0__tar.gz → 1.0.208.dev1__tar.gz

{code_loader-1.0.207.dev0 → code_loader-1.0.208.dev1}/PKG-INFO

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

{code_loader-1.0.207.dev0 → code_loader-1.0.208.dev1}/code_loader/contract/datasetclasses.py

@@ -6,13 +6,21 @@ import numpy as np
 import numpy.typing as npt
 
 from code_loader.contract.enums import DataStateType, DataStateEnum, LeapDataType, ConfusionMatrixValue, \
-    MetricDirection, DatasetMetadataType
+    MetricDirection, DatasetMetadataType, LatentSpaceReduction
 from code_loader.contract.visualizer_classes import LeapImage, LeapText, LeapGraph, LeapHorizontalBar, \
     LeapTextMask, LeapImageMask, LeapImageWithBBox, LeapImageWithHeatmap, LeapVideo, LeapAudio
 from code_loader.contract.sim_config import SimConfig
 
 custom_latent_space_attribute = "custom_latent_space"
 
+# Hard cap on a custom latent space's per-sample width, applied AFTER any reduction. The
+# soft threshold only warns. MAX_CUSTOM_LATENT_SPACES mirrors the engine's
+# src_tensorleap/common/types.py (source of truth) so an over-budget registration fails
+# locally instead of being silently dropped at eval time; keep the two in sync.
+MAX_CUSTOM_LATENT_SPACE_DIM = 4096
+CUSTOM_LATENT_SPACE_WARN_DIM = 1024
+MAX_CUSTOM_LATENT_SPACES = 10
+
 _simulation_context: Dict[str, bool] = {"active": False}
 
 SampleId = Union[int, str]
@@ -317,6 +325,11 @@ class CustomLatentSpaceHandler:
     name: str = 'custom_latent_space'
     use_ls_for_analysis: bool = False
     instance_aware: bool = False
+    computed_at: str = 'dataset'
+    arg_names: Optional[List[str]] = None
+    reduce: Optional[LatentSpaceReduction] = None
+    n_components: int = 512
+    channel_axis: int = -1
 
 
 # How a chain's latent-space vectors are derived from its per-step forward passes.

{code_loader-1.0.207.dev0 → code_loader-1.0.208.dev1}/code_loader/contract/enums.py

@@ -68,3 +68,8 @@ class ConfusionMatrixValue(Enum):
 class TestingSectionEnum(Enum):
     Warnings = "Warnings"
     Errors = "Errors"
+
+
+class LatentSpaceReduction(Enum):
+    MEAN_POOL = 'MEAN_POOL'
+    RANDOM_PROJECTION = 'RANDOM_PROJECTION'

{code_loader-1.0.207.dev0 → code_loader-1.0.208.dev1}/code_loader/contract/mapping.py

@@ -38,6 +38,7 @@ class NodeMappingType(Enum):
     Input8 = 'Input8'
     Input9 = 'Input9'
     PredictionLabels = 'PredictionLabels'
+    CustomLatentSpace = 'CustomLatentSpace'
 
 
 @dataclass

{code_loader-1.0.207.dev0 → code_loader-1.0.208.dev1}/code_loader/contract/responsedataclasses.py

@@ -144,6 +144,7 @@ class LeapAnalysisConfiguration:
 class EngineFileContract:
     node_connections: Optional[List[NodeConnection]] = None
     leap_analysis_configuration: Optional[LeapAnalysisConfiguration] = None
+    latent_space_connections: Optional[List[NodeConnection]] = None
 
 
 @dataclass

{code_loader-1.0.207.dev0 → code_loader-1.0.208.dev1}/code_loader/inner_leap_binder/leapbinder.py

@@ -17,12 +17,12 @@ from code_loader.contract.datasetclasses import SectionCallableInterface, InputH
     CustomMultipleReturnCallableInterfaceMultiArgs, DatasetBaseHandler, custom_latent_space_attribute, \
     RawInputsForHeatmap, VisualizerHandlerData, MetricHandlerData, CustomLossHandlerData, SamplePreprocessResponse, \
     ElementInstanceMasksHandler, InstanceCallableInterface, InstanceSectionCallableInterface, \
-    CustomLatentSpaceHandler, InstanceMetricHandler, \
+    CustomLatentSpaceHandler, InstanceMetricHandler, MAX_CUSTOM_LATENT_SPACES, \
     SimulationHandler, _simulation_context, AutoregressiveStepHandler, AutoregressiveStepCallableInterface, \
     AUTOREGRESSIVE_LATENT_SPACE_AGGREGATIONS, AUTOREGRESSIVE_IMPLICIT_ARG_NAMES, \
     AutoregressiveMetricHandler, AutoregressiveLossHandler, AutoregressiveVisualizerHandler
 from code_loader.contract.enums import LeapDataType, DataStateEnum, DataStateType, MetricDirection, DatasetMetadataType, \
-    TestingSectionEnum
+    TestingSectionEnum, LatentSpaceReduction
 from code_loader.contract.mapping import NodeConnection, NodeMapping, NodeMappingType
 from code_loader.contract.responsedataclasses import DatasetTestResultPayload, LeapAnalysisConfiguration
 from code_loader.contract.visualizer_classes import map_leap_data_type_to_visualizer_class
@@ -67,16 +67,29 @@ def _stringized_annotation_type_name(annotation: Any) -> Optional[str]:
     return None
 
 
-def _reject_stringized_sample_preprocess_response(function: Callable[..., Any], arg_name: str, annotation: Any) -> None:
-    """Fail fast on a stringized SamplePreprocessResponse arg (otherwise it silently
-    mis-wires as a regular input, surfacing only on the platform)."""
-    if _stringized_annotation_type_name(annotation) == SamplePreprocessResponse.__name__:
+def _reject_stringized_annotation(function: Callable[..., Any], arg_name: str, annotation: Any,
+                                  expected_cls: type, why: str) -> None:
+    """Fail fast on a stringized `expected_cls` arg (otherwise it silently mis-wires,
+    surfacing only on the platform)."""
+    if _stringized_annotation_type_name(annotation) == expected_cls.__name__:
         raise Exception(
             f"Argument '{arg_name}' of function '{function.__name__}' is annotated with a string "
-            f"('{annotation}') instead of the SamplePreprocessResponse type. This usually means the "
+            f"('{annotation}') instead of the {expected_cls.__name__} type. This usually means the "
             f"file uses 'from __future__ import annotations' (or a quoted hint), which stringizes "
-            f"annotations and breaks Tensorleap type detection. Remove that import (or the quotes) so "
-            f"SamplePreprocessResponse is referenced as a real type.")
+            f"annotations. {why} Remove that import (or the quotes).")
+
+
+def _reject_stringized_sample_preprocess_response(function: Callable[..., Any], arg_name: str, annotation: Any) -> None:
+    _reject_stringized_annotation(
+        function, arg_name, annotation, SamplePreprocessResponse,
+        "This breaks Tensorleap type detection, so it must be referenced as a real type.")
+
+
+def _reject_stringized_preprocess_response(function: Callable[..., Any], arg_name: str, annotation: Any) -> None:
+    _reject_stringized_annotation(
+        function, arg_name, annotation, PreprocessResponse,
+        "Tensorleap uses that annotation to tell a dataset-computed custom latent space from a "
+        "model-computed one, so it must be a real type.")
 
 
 
@@ -102,6 +115,7 @@ class LeapBinder:
         self._extend_with_default_losses()
 
         self.mapping_connections: List[NodeConnection] = []
+        self.latent_space_connections: List[NodeConnection] = []
         self.integration_test_func: Optional[Callable[[str, PreprocessResponse], Any]] = None
 
         self.batch_size_to_validate: Optional[int] = None
@@ -546,7 +560,12 @@ class LeapBinder:
     def set_custom_latent_space(self, function: Union[SectionCallableInterface, InstanceSectionCallableInterface],
                                 name: Optional[str] = None,
                                 use_ls_for_analysis: bool = False,
-                                instance_aware: bool = False) -> None:
+                                instance_aware: bool = False,
+                                computed_at: str = 'dataset',
+                                arg_names: Optional[List[str]] = None,
+                                reduce: Optional[LatentSpaceReduction] = None,
+                                n_components: int = 512,
+                                channel_axis: int = -1) -> None:
         """
         Register a custom latent space function.
 
@@ -567,7 +586,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. Not currently supported when
-            instance_aware=True — it is ignored (with a warning) and forced to False.
+            instance_aware=True, or for a model-computed latent space (computed_at='model')
+            — it is ignored (with a warning) and forced to False.
         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.
         """
@@ -579,6 +599,14 @@ class LeapBinder:
                 f"@tensorleap_custom_latent_space must have a unique name "
                 f"(pass name='...' to distinguish them)."
             )
+        if len(self.setup_container.custom_latent_spaces) >= MAX_CUSTOM_LATENT_SPACES:
+            raise Exception(
+                f"Cannot register custom latent space '{name}': Tensorleap supports at most "
+                f"{MAX_CUSTOM_LATENT_SPACES} custom latent spaces, and "
+                f"{len(self.setup_container.custom_latent_spaces)} are already registered "
+                f"({sorted(self.setup_container.custom_latent_spaces)}). Dataset-computed and "
+                f"model-computed latent spaces share this budget."
+            )
         # use_ls_for_analysis is not currently wired for instance-aware latent spaces (OOD /
         # Domain-Gap analyze the sample-level population, not instance rows) — force it off rather
         # than silently accepting a flag that has no effect.
@@ -588,6 +616,16 @@ class LeapBinder:
                 f"latent space ('{name}'). Ignoring it; the flag will be set to False."
             )
             use_ls_for_analysis = False
+        # Model-computed latent spaces are fetched via run_model_latent_space, a separate path
+        # that get_sample's per-sample custom_latent_spaces dict never populates — the engine has
+        # no consumer for a model-computed analysis latent space yet, so force it off rather than
+        # silently accepting a flag that has no effect.
+        if computed_at == 'model' and use_ls_for_analysis:
+            warnings.warn(
+                f"use_ls_for_analysis=True is not currently supported for a model-computed custom "
+                f"latent space ('{name}'). Ignoring it; the flag will be set to False."
+            )
+            use_ls_for_analysis = False
         if use_ls_for_analysis:
             already_flagged = [
                 existing_name
@@ -602,8 +640,25 @@ class LeapBinder:
                     f"Out-Of-Distribution and Domain-Gap insights). Set it on '{name}' "
                     f"or '{already_flagged[0]}', not both."
                 )
+        if reduce is not None:
+            if not isinstance(reduce, LatentSpaceReduction):
+                raise Exception(
+                    f"Custom latent space '{name}': reduce must be a LatentSpaceReduction, got "
+                    f"{type(reduce).__name__}.")
+            if reduce is LatentSpaceReduction.RANDOM_PROJECTION and (
+                    not isinstance(n_components, int) or isinstance(n_components, bool) or n_components <= 0):
+                raise Exception(
+                    f"Custom latent space '{name}': n_components must be a positive int for "
+                    f"LatentSpaceReduction.RANDOM_PROJECTION, got {n_components!r}.")
+            if reduce is LatentSpaceReduction.MEAN_POOL and (
+                    not isinstance(channel_axis, int) or isinstance(channel_axis, bool)):
+                raise Exception(
+                    f"Custom latent space '{name}': channel_axis must be an int for "
+                    f"LatentSpaceReduction.MEAN_POOL, got {channel_axis!r}.")
         self.setup_container.custom_latent_spaces[name] = CustomLatentSpaceHandler(
-            function=function, name=name, use_ls_for_analysis=use_ls_for_analysis, instance_aware=instance_aware)
+            function=function, name=name, use_ls_for_analysis=use_ls_for_analysis, instance_aware=instance_aware,
+            computed_at=computed_at, arg_names=arg_names, reduce=reduce, n_components=n_components,
+            channel_axis=channel_axis)
 
     def set_autoregressive_step(self, function: AutoregressiveStepCallableInterface,
                                 latent_space_aggregation: str = 'last_step') -> None:

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

@@ -26,15 +26,17 @@ from code_loader.contract.datasetclasses import CustomCallableInterfaceMultiArgs
     CustomMultipleReturnCallableInterfaceMultiArgs, ConfusionMatrixCallableInterfaceMultiArgs, CustomCallableInterface, \
     VisualizerCallableInterface, MetadataSectionCallableInterface, PreprocessResponse, SectionCallableInterface, \
     ConfusionMatrixElement, SamplePreprocessResponse, PredictionTypeHandler, InstanceCallableInterface, ElementInstance, \
-    InstanceLengthCallableInterface, InstanceSectionCallableInterface, AutoregressiveStepCallableInterface
+    InstanceLengthCallableInterface, InstanceSectionCallableInterface, AutoregressiveStepCallableInterface, \
+    MAX_CUSTOM_LATENT_SPACE_DIM, CUSTOM_LATENT_SPACE_WARN_DIM
 from code_loader.contract.enums import MetricDirection, LeapDataType, DatasetMetadataType, DataStateType, \
-    DataStateEnum
+    DataStateEnum, LatentSpaceReduction
 from code_loader import leap_binder, LeapLoader
 from code_loader.contract.mapping import NodeMapping, NodeMappingType, NodeConnection
 from code_loader.contract.visualizer_classes import LeapImage, LeapImageMask, LeapTextMask, LeapText, LeapGraph, \
     LeapHorizontalBar, LeapImageWithBBox, LeapImageWithHeatmap, LeapVideo, LeapAudio, LeapValidationError, \
     map_leap_data_type_to_visualizer_class
-from code_loader.inner_leap_binder.leapbinder import mapping_runtime_mode_env_var_mame
+from code_loader.inner_leap_binder.leapbinder import mapping_runtime_mode_env_var_mame, \
+    _reject_stringized_preprocess_response, _reject_stringized_sample_preprocess_response
 from code_loader.mixpanel_tracker import clear_integration_events, AnalyticsEvent, emit_integration_event_once
 
 _called_from_inside_tl_decorator = 0
@@ -264,7 +266,8 @@ def _validate_grouped_result(result, group_size, func_name, validate_single):
             f'{group_size} arrays, got {type(result)}.')
 
 
-def _add_mapping_connection(user_unique_name, connection_destinations, arg_names, name, node_mapping_type):
+def _add_mapping_connection(user_unique_name, connection_destinations, arg_names, name, node_mapping_type,
+                            target_list=None):
     connection_destinations = [connection_destination for connection_destination in connection_destinations
                                if not isinstance(connection_destination, SamplePreprocessResponse)]
 
@@ -274,7 +277,9 @@ def _add_mapping_connection(user_unique_name, connection_destinations, arg_names
     for arg_name, destination in zip(arg_names, connection_destinations):
         node_inputs[arg_name] = destination.node_mapping
 
-    leap_binder.mapping_connections.append(NodeConnection(main_node_mapping, node_inputs))
+    if target_list is None:
+        target_list = leap_binder.mapping_connections
+    target_list.append(NodeConnection(main_node_mapping, node_inputs))
 
 
 def _add_mapping_connections(connects_to, arg_names, node_mapping_type, name):
@@ -1750,7 +1755,92 @@ def tensorleap_metadata(
     return decorating_function
 
 
-def tensorleap_custom_latent_space(name: Optional[str] = None, use_ls_for_analysis: bool = False):
+def _classify_custom_latent_space_signature(user_function) -> str:
+    argspec = inspect.getfullargspec(user_function)
+    params = argspec.args
+    if len(params) != 2:
+        return 'model'
+
+    second = params[1]
+    first = params[0]
+    if second not in argspec.annotations:
+        raise Exception(
+            f"tensorleap_custom_latent_space validation failed: '{user_function.__name__}' has "
+            f"exactly two parameters ('{first}', '{second}') and '{second}' has no type "
+            f"annotation, so Tensorleap cannot tell whether this is a dataset-computed latent "
+            f"space (one sample at a time) or a model-computed one (a batch of model tensors). "
+            f"Please annotate the second parameter '{second}' to disambiguate:\n"
+            f"  dataset-computed: def {user_function.__name__}({first}, {second}: PreprocessResponse) -> (d,)\n"
+            f"  model-computed:   def {user_function.__name__}({first}: np.ndarray, {second}: np.ndarray) -> (batch, d)\n"
+            f"If you are upgrading an existing project, this signature used to be accepted "
+            f"unannotated as dataset-computed; add ': PreprocessResponse' to '{second}' to keep "
+            f"the previous behavior.")
+
+    annotation = argspec.annotations[second]
+    _reject_stringized_preprocess_response(user_function, second, annotation)
+    if annotation is PreprocessResponse:
+        return 'dataset'
+    return 'model'
+
+
+def _model_latent_space_arg_names(user_function) -> List[str]:
+    argspec = inspect.getfullargspec(user_function)
+    arg_names = list(argspec.args)
+    preprocess_response_arg_name = None
+    for arg_name, arg_type in argspec.annotations.items():
+        _reject_stringized_sample_preprocess_response(user_function, arg_name, arg_type)
+        if arg_type == SamplePreprocessResponse:
+            if preprocess_response_arg_name is not None:
+                raise Exception(
+                    f"tensorleap_custom_latent_space validation failed: only one argument of "
+                    f"'{user_function.__name__}' can be of type SamplePreprocessResponse.")
+            preprocess_response_arg_name = arg_name
+            arg_names.remove(arg_name)
+    return arg_names
+
+
+def _custom_latent_space_width(result, ls_name, reduce, n_components, channel_axis,
+                               has_batch_axis: bool) -> int:
+    if reduce is LatentSpaceReduction.RANDOM_PROJECTION:
+        return n_components
+    if reduce is LatentSpaceReduction.MEAN_POOL:
+        axis = channel_axis if channel_axis >= 0 else result.ndim + channel_axis
+        min_axis = 1 if has_batch_axis else 0
+        if axis < min_axis or axis >= result.ndim:
+            raise Exception(
+                f"tensorleap_custom_latent_space validation failed: '{ls_name}' uses "
+                f"LatentSpaceReduction.MEAN_POOL with channel_axis={channel_axis}, which does not "
+                f"select a non-batch axis of the returned shape {tuple(result.shape)}."
+                + (" Axis 0 is the batch." if has_batch_axis else ""))
+        return int(result.shape[axis])
+    dims = result.shape[1:] if has_batch_axis else result.shape
+    return int(np.prod(dims))
+
+
+def _check_custom_latent_space_width(result, ls_name, reduce, n_components, channel_axis,
+                                     has_batch_axis: bool) -> None:
+    width = _custom_latent_space_width(result, ls_name, reduce, n_components, channel_axis,
+                                       has_batch_axis)
+    if width > MAX_CUSTOM_LATENT_SPACE_DIM:
+        raise Exception(
+            f"tensorleap_custom_latent_space validation failed: '{ls_name}' produces "
+            f"{width} dimensions per sample, above the {MAX_CUSTOM_LATENT_SPACE_DIM} limit. "
+            f"Pass reduce=LatentSpaceReduction.RANDOM_PROJECTION (with n_components) to "
+            f"project it down, or reduce=LatentSpaceReduction.MEAN_POOL (with channel_axis) "
+            f"to average the non-channel axes, or return a smaller array.")
+    if width > CUSTOM_LATENT_SPACE_WARN_DIM:
+        store_general_warning(
+            key=("tensorleap_custom_latent_space_width", ls_name, width),
+            message=(
+                f"Custom latent space '{ls_name}' produces {width} dimensions per sample. "
+                f"Wide latent spaces are slow to transit and store. Consider "
+                f"reduce=LatentSpaceReduction.MEAN_POOL or "
+                f"reduce=LatentSpaceReduction.RANDOM_PROJECTION."))
+
+
+def tensorleap_custom_latent_space(name: Optional[str] = None, use_ls_for_analysis: bool = False,
+                                   reduce: Optional[LatentSpaceReduction] = None,
+                                   n_components: int = 512, channel_axis: int = -1):
     assert isinstance(use_ls_for_analysis, bool), \
         ("tensorleap_custom_latent_space validation failed: use_ls_for_analysis must be a bool. "
          f"Got {type(use_ls_for_analysis)}.")
@@ -1758,6 +1848,10 @@ def tensorleap_custom_latent_space(name: Optional[str] = None, use_ls_for_analys
     def decorating_function(user_function: SectionCallableInterface):
         ls_name = name if name is not None else user_function.__name__
 
+        if _classify_custom_latent_space_signature(user_function) == 'model':
+            return _decorate_model_latent_space(user_function, ls_name, use_ls_for_analysis,
+                                                reduce, n_components, channel_axis)
+
         def _validate_input_args(sample_id: Union[int, str, list], preprocess_response: PreprocessResponse):
             _validate_id_or_group(sample_id, preprocess_response, 'tensorleap_custom_latent_space')
 
@@ -1777,6 +1871,8 @@ def tensorleap_custom_latent_space(name: Optional[str] = None, use_ls_for_analys
                         f"inside your function."
                     ),
                 )
+            _check_custom_latent_space_width(single_result, ls_name, reduce, n_components, channel_axis,
+                                             has_batch_axis=False)
 
         def _validate_result(result, grouped=False, group_size=None):
             if not grouped:
@@ -1814,7 +1910,9 @@ def tensorleap_custom_latent_space(name: Optional[str] = None, use_ls_for_analys
             return result
 
         leap_binder.set_custom_latent_space(inner_without_validate, ls_name,
-                                            use_ls_for_analysis=use_ls_for_analysis)
+                                            use_ls_for_analysis=use_ls_for_analysis,
+                                            computed_at='dataset', reduce=reduce,
+                                            n_components=n_components, channel_axis=channel_axis)
 
         def inner(sample_id, preprocess_response):
             if os.environ.get(mapping_runtime_mode_env_var_mame):
@@ -1834,6 +1932,92 @@ def tensorleap_custom_latent_space(name: Optional[str] = None, use_ls_for_analys
     return decorating_function
 
 
+def _decorate_model_latent_space(user_function, ls_name, use_ls_for_analysis, reduce, n_components,
+                                channel_axis):
+    arg_names = _model_latent_space_arg_names(user_function)
+
+    def _validate_input_args(*args, **kwargs):
+        assert len(args) + len(kwargs) > 0, (
+            f"tensorleap_custom_latent_space validation failed: '{user_function.__name__}' is "
+            f"model-computed and expects at least one np.ndarray argument, but received none.")
+        for i, arg in enumerate(args):
+            assert isinstance(arg, (np.ndarray, SamplePreprocessResponse)), (
+                f"tensorleap_custom_latent_space validation failed: Argument #{i} of "
+                f"'{user_function.__name__}' should be a numpy array. Got {type(arg)}.")
+        for arg_name, arg in kwargs.items():
+            assert isinstance(arg, (np.ndarray, SamplePreprocessResponse)), (
+                f"tensorleap_custom_latent_space validation failed: Argument {arg_name} of "
+                f"'{user_function.__name__}' should be a numpy array. Got {type(arg)}.")
+
+    def _validate_result(result):
+        assert isinstance(result, np.ndarray), (
+            f"tensorleap_custom_latent_space validation failed: '{user_function.__name__}' is "
+            f"model-computed and should return a numpy array of shape (batch, d). "
+            f"Got {type(result)}.")
+        assert result.ndim >= 2, (
+            f"tensorleap_custom_latent_space validation failed: '{user_function.__name__}' "
+            f"returned shape {tuple(result.shape)}. A model-computed latent space returns "
+            f"(batch, d), so the result needs a batch axis and at least one feature axis.")
+        if leap_binder.batch_size_to_validate:
+            assert result.shape[0] == leap_binder.batch_size_to_validate, (
+                f"tensorleap_custom_latent_space validation failed: '{user_function.__name__}' "
+                f"returned leading dim {result.shape[0]} instead of the batch size "
+                f"{leap_binder.batch_size_to_validate}.")
+
+        _check_custom_latent_space_width(result, ls_name, reduce, n_components, channel_axis,
+                                         has_batch_axis=True)
+
+    def inner_without_validate(*args, **kwargs):
+        global _called_from_inside_tl_decorator
+        _called_from_inside_tl_decorator += 1
+
+        try:
+            result = user_function(*args, **kwargs)
+        finally:
+            _called_from_inside_tl_decorator -= 1
+
+        return result
+
+    inner_without_validate.__signature__ = inspect.signature(user_function)
+
+    leap_binder.set_custom_latent_space(inner_without_validate, ls_name,
+                                        use_ls_for_analysis=use_ls_for_analysis,
+                                        computed_at='model', arg_names=arg_names, reduce=reduce,
+                                        n_components=n_components, channel_axis=channel_axis)
+
+    def inner(*args, **kwargs):
+        args, kwargs, _was_downcast = _simulate_engine_tensor_dtype(args, kwargs)
+        _validate_input_args(*args, **kwargs)
+        result = inner_without_validate(*args, **kwargs)
+        _validate_result(result)
+        return result
+
+    def mapping_inner(*args, **kwargs):
+        if _mapping_dataset_is_grouped:
+            raise LeapValidationError(
+                f"tensorleap_custom_latent_space validation failed: '{ls_name}' is model-computed, "
+                f"which is not supported for a grouped preprocess response. Use the "
+                f"(sample_id, preprocess: PreprocessResponse) form instead.")
+
+        ordered_connections = [kwargs[arg_name] for arg_name in arg_names if arg_name in kwargs]
+        ordered_connections = list(args) + ordered_connections
+
+        leap_binder.latent_space_connections[:] = [
+            connection for connection in leap_binder.latent_space_connections
+            if connection.node.name != ls_name]
+        _add_mapping_connection(ls_name, ordered_connections, arg_names, ls_name,
+                                NodeMappingType.CustomLatentSpace,
+                                target_list=leap_binder.latent_space_connections)
+        return None
+
+    def final_inner(*args, **kwargs):
+        if os.environ.get(mapping_runtime_mode_env_var_mame):
+            return mapping_inner(*args, **kwargs)
+        return inner(*args, **kwargs)
+
+    return final_inner
+
+
 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. "

{code_loader-1.0.207.dev0 → code_loader-1.0.208.dev1}/code_loader/leaploader.py

@@ -352,6 +352,9 @@ class LeapLoader(LeapLoaderBase):
                 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)
+                model_ls_test_payload = self._check_model_latent_spaces()
+                if model_ls_test_payload is not None:
+                    test_payloads.append(model_ls_test_payload)
                 handlers_test_payloads = self._check_handlers()
                 test_payloads.extend(handlers_test_payloads)
                 simulation_test_payloads = self._check_simulations()
@@ -386,8 +389,10 @@ class LeapLoader(LeapLoaderBase):
                                        is_valid_for_model=is_valid_for_model, setup=setup_response,
                                        model_setup=model_setup, general_error=general_error,
                                        print_log=print_log,
-                                       engine_file_contract=EngineFileContract(global_leap_binder.mapping_connections,
-                                                                               global_leap_binder.leap_analysis_configuration))
+                                       engine_file_contract=EngineFileContract(
+                                           global_leap_binder.mapping_connections,
+                                           global_leap_binder.leap_analysis_configuration,
+                                           global_leap_binder.latent_space_connections))
 
     def _check_integration_test_exists(self) -> DatasetTestResultPayload:
         test_result = DatasetTestResultPayload('integration_test')
@@ -425,6 +430,28 @@ class LeapLoader(LeapLoaderBase):
             test_result.is_passed = False
         return test_result
 
+    def _check_model_latent_spaces(self) -> Optional[DatasetTestResultPayload]:
+        model_names = [
+            name for name, handler in global_leap_binder.setup_container.custom_latent_spaces.items()
+            if handler.computed_at == 'model'
+        ]
+        if not model_names:
+            return None
+
+        test_result = DatasetTestResultPayload('model_custom_latent_space')
+        grouped_states = [
+            state.name for state, preprocess_response in self._preprocess_result().items()
+            if preprocess_response.is_grouped
+        ]
+        if grouped_states:
+            test_result.is_passed = False
+            test_result.display[TestingSectionEnum.Errors.name] = (
+                f"Model-computed custom latent space(s) {model_names} are not supported with a "
+                f"grouped preprocess response (grouped: {grouped_states}). Use the "
+                f"(sample_id, preprocess: PreprocessResponse) form instead."
+            )
+        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()
@@ -1390,7 +1417,9 @@ class LeapLoader(LeapLoaderBase):
             sample_id: Union[int, str],
             preprocess: "PreprocessResponse",
             instance_id: Optional[int] = None) -> Optional[Dict[str, npt.NDArray[np.float32]]]:
-        handlers = global_leap_binder.setup_container.custom_latent_spaces
+        handlers = {handler_name: handler for handler_name, handler
+                    in global_leap_binder.setup_container.custom_latent_spaces.items()
+                    if handler.computed_at != 'model'}
         if not handlers:
             return None
         if preprocess.is_grouped:
@@ -1428,6 +1457,46 @@ class LeapLoader(LeapLoaderBase):
         self.exec_script()
         return tuple(global_leap_binder.setup_container.custom_latent_spaces.keys())
 
+    @lru_cache()
+    def get_dataset_custom_latent_space_names(self) -> Tuple[str, ...]:
+        self.exec_script()
+        return tuple(name for name, handler in global_leap_binder.setup_container.custom_latent_spaces.items()
+                     if handler.computed_at != 'model')
+
+    @lru_cache()
+    def get_model_custom_latent_space_names(self) -> Tuple[str, ...]:
+        self.exec_script()
+        return tuple(name for name, handler in global_leap_binder.setup_container.custom_latent_spaces.items()
+                     if handler.computed_at == 'model')
+
+    @lru_cache()
+    def get_model_latent_space_specs(self) -> Dict[str, Dict[str, Any]]:
+        self.exec_script()
+        return {
+            name: {
+                'arg_names': list(handler.arg_names or []),
+                'reduce': handler.reduce.value if handler.reduce is not None else None,
+                'n_components': handler.n_components,
+                'channel_axis': handler.channel_axis,
+            }
+            for name, handler in global_leap_binder.setup_container.custom_latent_spaces.items()
+            if handler.computed_at == 'model'
+        }
+
+    def run_model_latent_space(self, ls_name: str, sample_ids: np.array, state: DataStateEnum,
+                               input_tensors_by_arg_name: Dict[str, npt.NDArray[np.float32]]
+                               ) -> npt.NDArray[np.float32]:
+        self._preprocess_result()
+
+        handler = global_leap_binder.setup_container.custom_latent_spaces[ls_name]
+        preprocess_response_arg_name = self._get_preprocess_response_arg_name(handler.function)
+
+        if preprocess_response_arg_name is not None:
+            input_tensors_by_arg_name[preprocess_response_arg_name] = SamplePreprocessResponse(
+                sample_ids, self._preprocess_result()[state])
+
+        return handler.function(**input_tensors_by_arg_name)
+
     @lru_cache()
     def get_instance_custom_latent_space_names(self) -> Tuple[str, ...]:
         """Names of registered custom latent spaces that are instance-aware.

{code_loader-1.0.207.dev0 → code_loader-1.0.208.dev1}/code_loader/leaploaderbase.py

@@ -226,6 +226,31 @@ class LeapLoaderBase:
     def get_custom_latent_space_for_analysis(self) -> Optional[str]:
         pass
 
+    # These raise rather than `pass` for the same reason as the autoregressive entry points
+    # above: an un-overridden `pass` body would report "no model-computed latent spaces" on a
+    # loader that simply predates them, and the engine would silently skip computing them.
+    @abstractmethod
+    def get_dataset_custom_latent_space_names(self) -> Tuple[str, ...]:
+        raise NotImplementedError(f'{type(self).__name__} does not implement '
+                                  'get_dataset_custom_latent_space_names.')
+
+    @abstractmethod
+    def get_model_custom_latent_space_names(self) -> Tuple[str, ...]:
+        raise NotImplementedError(f'{type(self).__name__} does not implement '
+                                  'get_model_custom_latent_space_names.')
+
+    @abstractmethod
+    def get_model_latent_space_specs(self) -> Dict[str, Dict[str, Any]]:
+        raise NotImplementedError(f'{type(self).__name__} does not implement '
+                                  'get_model_latent_space_specs.')
+
+    @abstractmethod
+    def run_model_latent_space(self, ls_name: str, sample_ids: np.array, state: DataStateEnum,
+                               input_tensors_by_arg_name: Dict[str, npt.NDArray[np.float32]]
+                               ) -> npt.NDArray[np.float32]:
+        raise NotImplementedError(f'{type(self).__name__} does not implement '
+                                  'run_model_latent_space.')
+
     @abstractmethod
     def get_heatmap_visualizer_raw_vis_input_arg_name(self, visualizer_name: str) -> Optional[str]:
         pass

{code_loader-1.0.207.dev0 → code_loader-1.0.208.dev1}/pyproject.toml

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