code-loader 1.0.204__tar.gz → 1.0.205.dev0__tar.gz

{code_loader-1.0.204 → code_loader-1.0.205.dev0}/PKG-INFO

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

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

@@ -1032,17 +1032,18 @@ 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 calls
-        because it consults setup_container.inputs and prediction_types.
+        Order-independent: this runs after all set_input / add_prediction /
+        set_custom_latent_space calls because it consults setup_container.
         """
         if not self.leap_analysis_configuration:
             return
@@ -1050,24 +1051,37 @@ class LeapBinder:
         if not names:
             return
 
-        forbidden = {"balanced", "user_custom"}
+        forbidden = {"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"} | input_names | prediction_names
+        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
+        ]
 
         for name in names:
             if name in forbidden:
                 raise Exception(
-                    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."
+                    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."
                 )
             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']."
+                    f"tokens=['foreground', 'balanced']."
                 )
 
     def set_batch_size_to_validate(self, batch_size: int) -> None:

{code_loader-1.0.204 → code_loader-1.0.205.dev0}/code_loader/inner_leap_binder/leapbinder_decorators.py

@@ -16,7 +16,9 @@ import numpy as np
 import numpy.typing as npt
 
 from code_loader.utils import get_metadata_type_from_variable, map_dict_to_metadata_types, \
-    validate_autoregressive_state_types, autoregressive_nests_equal
+    validate_autoregressive_state_types, autoregressive_nests_equal, \
+    simulate_engine_float16_downcast_on_call_args, ENGINE_STORAGE_DTYPE, \
+    TL_DISABLE_ENGINE_FLOAT16_SIMULATION_ENV_VAR
 
 logger = logging.getLogger(__name__)
 
@@ -308,6 +310,35 @@ def _require_sample_preprocess_response_supplied(user_function: Callable, args:
             f"SamplePreprocessResponse(sample_id, preprocess)).")
 
 
+def _simulate_engine_tensor_dtype(args: tuple, kwargs: dict):
+    """Hand a metric / loss the tensor dtype the platform will hand it.
+
+    The engine downcasts every tensor to float16 before pushing it to Redis, and metrics and
+    losses read their tensors back from there. A local run would otherwise pass the float32
+    arrays straight from the encoders and the model, so a float32-only loss (a Keras/TF loss
+    that refuses a float16 y_pred, say) passes the local test and only fails in evaluation.
+    Downcasting here makes the local test fail the same way, before the project is pushed.
+
+    No-op on the platform, where the tensors already arrive downcast. Returns the (possibly
+    rewritten) call args plus whether anything was actually downcast.
+    """
+    if _call_from_tl_platform:
+        return args, kwargs, False
+    downcast_args, downcast_kwargs = simulate_engine_float16_downcast_on_call_args(args, kwargs)
+    was_downcast = any(arg is not original for arg, original in zip(downcast_args, args)) or \
+        any(arg is not kwargs[arg_name] for arg_name, arg in downcast_kwargs.items())
+    return downcast_args, downcast_kwargs, was_downcast
+
+
+def _print_engine_float16_simulation_hint(func_name: str) -> None:
+    print(f"note: {func_name}() received its tensors as {np.dtype(ENGINE_STORAGE_DTYPE).name} — "
+          f"the platform stores every tensor at that precision, so the local test feeds the same "
+          f"dtype the engine will. If the failure above is a dtype mismatch, cast inside "
+          f"{func_name}() (e.g. arr.astype(np.float32)) rather than relying on the caller's "
+          f"dtype. Set {TL_DISABLE_ENGINE_FLOAT16_SIMULATION_ENV_VAR}=true to opt out of this "
+          f"simulation locally.")
+
+
 def _warn_loss_inputs_not_wired_to_model() -> None:
     """Warn when a custom loss consumes an input encoder the model is not given.
 
@@ -900,7 +931,10 @@ def tensorleap_custom_metric(name: str,
                         )
 
                     else:
-                        assert all(isinstance(v, (float, int, type(None), np.float32)) for v in single_metric_result), (
+                        # np.floating, not just np.float32: metric inputs are downcast to the
+                        # engine's float16, so a metric that returns its computed values as numpy
+                        # scalars now yields np.float16 ones.
+                        assert all(isinstance(v, (float, int, type(None), np.floating)) for v in single_metric_result), (
                             f"{supported_types_message}\n"
                             f"Got {'a dict where the value of ' + str(key) + ' is of type ' if key is not None else ''}"
                             f"List[{', '.join(type(v).__name__ for v in single_metric_result)}]."
@@ -1019,9 +1053,15 @@ def tensorleap_custom_metric(name: str,
         def inner(*args, **kwargs):
             if not _call_from_tl_platform:
                 set_current('tensorleap_custom_metric')
+            args, kwargs, was_downcast = _simulate_engine_tensor_dtype(args, kwargs)
             _validate_input_args(*args, **kwargs)
 
-            result = inner_without_validate(*args, **kwargs)
+            try:
+                result = inner_without_validate(*args, **kwargs)
+            except Exception:
+                if was_downcast:
+                    _print_engine_float16_simulation_hint(user_function.__name__)
+                raise
 
             _validate_result(result)
             if not _call_from_tl_platform:
@@ -1229,9 +1269,15 @@ def tensorleap_custom_instances_metric(name: str,
         def inner(*args, **kwargs):
             if not _call_from_tl_platform:
                 set_current('tensorleap_custom_instances_metric')
+            args, kwargs, was_downcast = _simulate_engine_tensor_dtype(args, kwargs)
             _validate_input_args(*args, **kwargs)
 
-            result = inner_without_validate(*args, **kwargs)
+            try:
+                result = inner_without_validate(*args, **kwargs)
+            except Exception:
+                if was_downcast:
+                    _print_engine_float16_simulation_hint(user_function.__name__)
+                raise
 
             _validate_result(result)
             if not _call_from_tl_platform:
@@ -3454,9 +3500,15 @@ def tensorleap_custom_loss(name: str, connects_to=None):
         def inner(*args, **kwargs):
             if not _call_from_tl_platform:
                 set_current("tensorleap_custom_loss")
+            args, kwargs, was_downcast = _simulate_engine_tensor_dtype(args, kwargs)
             _validate_input_args(*args, **kwargs)
 
-            result = inner_without_validate(*args, **kwargs)
+            try:
+                result = inner_without_validate(*args, **kwargs)
+            except Exception:
+                if was_downcast:
+                    _print_engine_float16_simulation_hint(user_function.__name__)
+                raise
 
             _validate_result(result)
             if not _call_from_tl_platform:

{code_loader-1.0.204 → code_loader-1.0.205.dev0}/code_loader/utils.py

@@ -1,5 +1,6 @@
 import io
 import math
+import os
 import pickle
 import sys
 from pathlib import Path
@@ -42,6 +43,80 @@ def to_numpy_return_masks_wrapper(encoder_function: InstanceCallableInterface) -
     return numpy_encoder_function
 
 
+# The engine downcasts every tensor to float16 before pushing it to Redis, and metrics and losses
+# read their tensors back from there — so on the platform they are handed float16 arrays. A local
+# run feeds them the float32 arrays straight from the encoders and the model, which is why a
+# float32-only loss (e.g. a Keras/TF loss that refuses a float16 y_pred) passes the local test and
+# then fails in evaluation. Mirror the engine locally so the failure surfaces before the push.
+ENGINE_STORAGE_DTYPE = np.float16
+# The engine's cast is range-guarded (redis_metrics_queue_manager.py): only a non-empty float32 /
+# float64 array whose abs-max fits float16's finite range is cast, so an array holding large values
+# (big logits, say) keeps its original dtype instead of overflowing to inf. Arrays holding NaN are
+# left alone too — their abs-max is NaN, and `NaN <= max` is False.
+ENGINE_STORAGE_DTYPE_MAX = 65504.0  # float16 finite max
+ENGINE_DOWNCAST_SOURCE_DTYPES = (np.float32, np.float64)
+TL_DISABLE_ENGINE_FLOAT16_SIMULATION_ENV_VAR = 'TL_DISABLE_ENGINE_FLOAT16_SIMULATION'
+
+
+def is_engine_float16_simulation_disabled() -> bool:
+    return os.environ.get(TL_DISABLE_ENGINE_FLOAT16_SIMULATION_ENV_VAR, '').lower() in \
+        ('1', 'true', 'yes')
+
+
+def simulate_engine_float16_downcast(value: Any) -> Any:
+    """Return ``value`` with its float tensors downcast the way the engine downcasts them.
+
+    Recurses into the containers a handler argument can arrive in (a grouped input is a list of
+    arrays). Only the arrays the engine would actually cast are touched — see
+    :func:`_downcast_array_to_engine_dtype`. Ints, bools and strings survive the engine's storage
+    unchanged, and anything that is not a tensor or a container of tensors
+    (SamplePreprocessResponse, for instance) is returned as is.
+    """
+    if is_engine_float16_simulation_disabled():
+        return value
+    return _downcast_to_engine_dtype(value)
+
+
+def _downcast_array_to_engine_dtype(array: npt.NDArray[Any]) -> npt.NDArray[Any]:
+    """Mirror the engine's guarded cast for a single array.
+
+    The engine casts an array to float16 only when it is a non-empty float32 / float64 array whose
+    whole-array abs-max fits float16's finite range; anything else is stored at its original dtype.
+    Casting unconditionally would diverge the other way — a local run would hand a handler ``inf``
+    where the platform hands it a finite float32.
+    """
+    if array.dtype not in ENGINE_DOWNCAST_SOURCE_DTYPES or not array.size:
+        return array
+    # NaN fails this comparison, which is how the engine ends up leaving NaN-bearing arrays alone.
+    if not float(np.abs(array).max()) <= ENGINE_STORAGE_DTYPE_MAX:
+        return array
+    return array.astype(ENGINE_STORAGE_DTYPE)
+
+
+def _downcast_to_engine_dtype(value: Any) -> Any:
+    # Containers return the original object when no leaf changed, so callers can tell whether a
+    # downcast actually happened by identity.
+    if isinstance(value, np.ndarray):
+        return _downcast_array_to_engine_dtype(value)
+    if isinstance(value, (list, tuple)):
+        downcast = [_downcast_to_engine_dtype(item) for item in value]
+        if all(item is original for item, original in zip(downcast, value)):
+            return value
+        return downcast if isinstance(value, list) else tuple(downcast)
+    if isinstance(value, dict):
+        downcast_dict = {key: _downcast_to_engine_dtype(item) for key, item in value.items()}
+        if all(item is value[key] for key, item in downcast_dict.items()):
+            return value
+        return downcast_dict
+    return value
+
+
+def simulate_engine_float16_downcast_on_call_args(
+        args: Tuple[Any, ...], kwargs: Dict[str, Any]) -> Tuple[Tuple[Any, ...], Dict[str, Any]]:
+    return (tuple(simulate_engine_float16_downcast(arg) for arg in args),
+            {arg_name: simulate_engine_float16_downcast(arg) for arg_name, arg in kwargs.items()})
+
+
 def get_root_traceback(exc_tb: TracebackType) -> TracebackType:
     return_traceback = exc_tb
     while return_traceback.tb_next is not None:

{code_loader-1.0.204 → code_loader-1.0.205.dev0}/pyproject.toml

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