code-loader 1.0.197.dev0__py3-none-any.whl → 1.0.197.dev2__py3-none-any.whl

code_loader/inner_leap_binder/leapbinder_decorators.py

@@ -40,6 +40,7 @@ _called_from_inside_tl_integration_test_decorator = False
 # to tell a raise in the test body (report as code-flow failure) from a module-level crash
 # before any test ran (report as a plain script crash).
 _integration_test_started = False
+_mapping_dataset_is_grouped = False
 _call_from_tl_platform = os.environ.get('IS_TENSORLEAP_PLATFORM') == 'true'
 
 # ---- warnings store (module-level) ----
@@ -213,6 +214,53 @@ def batch_warning(result, func_name):
         )
 
 
+def _validate_id_or_group(sample_id, preprocess_response, func_name):
+    """Accept either a single sample id (flat dataset) or a group — a list of ids (grouped
+    dataset). The argument shape must match the dataset: a group requires a grouped
+    PreprocessResponse and a single id requires a flat one. Every id must match the declared
+    sample_id_type."""
+    is_group = isinstance(sample_id, list)
+    if is_group and not preprocess_response.is_grouped:
+        raise AssertionError(
+            f'{func_name}() validation failed: got a group (list of sample ids) but the '
+            f'PreprocessResponse is not grouped. Pass a single sample id for a flat dataset.')
+    if not is_group and preprocess_response.is_grouped:
+        raise AssertionError(
+            f'{func_name}() validation failed: got a single sample id but the PreprocessResponse '
+            f'is grouped. Pass a group (list of sample ids) for a grouped dataset.')
+    ids = sample_id if is_group else [sample_id]
+    for sid in ids:
+        assert type(sid) == preprocess_response.sample_id_type, \
+            (f'{func_name}() validation failed: '
+             f'Argument sample_id should be as the same type as defined in the preprocess response '
+             f'{preprocess_response.sample_id_type}. Got {type(sid)}.')
+
+
+def _validate_grouped_result(result, group_size, func_name, validate_single):
+    """A grouped (group) encoder result must be either a single ndarray with leading dim
+    == group_size, or a list of `group_size` per-sample results (each validated by
+    validate_single). Metadata validates its grouped result separately (list only).
+    An `(B, *)` array is validated per-sample slice (not as the whole batched array) so
+    per-sample checks (e.g. channel_dim <= rank) use the sample shape, not the batch shape."""
+    if isinstance(result, np.ndarray):
+        assert len(result.shape) > 0 and result.shape[0] == group_size, \
+            (f'{func_name}() validation failed: expected a grouped output for a group of '
+             f'{group_size}: an array with leading dim {group_size}, got shape {result.shape}.')
+        for single in result:
+            validate_single(single)
+    elif isinstance(result, list):
+        assert len(result) == group_size, \
+            (f'{func_name}() validation failed: expected a grouped output for a group of '
+             f'{group_size}: a list of {group_size} arrays, got {len(result)}.')
+        for r in result:
+            validate_single(r)
+    else:
+        raise AssertionError(
+            f'{func_name}() validation failed: expected a grouped output for a group of '
+            f'{group_size}: either an array with leading dim {group_size} or a list of '
+            f'{group_size} arrays, got {type(result)}.')
+
+
 def _add_mapping_connection(user_unique_name, connection_destinations, arg_names, name, node_mapping_type):
     connection_destinations = [connection_destination for connection_destination in connection_destinations
                                if not isinstance(connection_destination, SamplePreprocessResponse)]
@@ -290,11 +338,7 @@ def tensorleap_integration_test():
 
         def _validate_input_args(*args, **kwargs):
             sample_id, preprocess_response = args
-            assert type(sample_id) == preprocess_response.sample_id_type, (
-                f"tensorleap_integration_test validation failed: "
-                f"sample_id type ({type(sample_id).__name__}) does not match the expected "
-                f"type ({preprocess_response.sample_id_type}) from the PreprocessResponse."
-            )
+            _validate_id_or_group(sample_id, preprocess_response, 'tensorleap_integration_test')
 
         def inner(*args, **kwargs):
             if not _call_from_tl_platform:
@@ -305,7 +349,7 @@ def tensorleap_integration_test():
             expected_names = inspect.getfullargspec(integration_test_function)[0][:2]
             if len(expected_names) < 2:
                 expected_names = ['idx', 'preprocess']
-            validate_args_structure(*args, types_order=[Union[int, str], PreprocessResponse],
+            validate_args_structure(*args, types_order=[Union[int, str, list], PreprocessResponse],
                                     func_name='integration_test', expected_names=expected_names,
                                     **kwargs)
             sample_id, preprocess_response = (
@@ -329,6 +373,8 @@ def tensorleap_integration_test():
                 _reset_model_loop_state()
                 ret = integration_test_function(*args, **kwargs)
 
+                global _mapping_dataset_is_grouped
+                _mapping_dataset_is_grouped = args[1].is_grouped
                 try:
                     os.environ[mapping_runtime_mode_env_var_mame] = 'True'
                     _reset_model_loop_state()
@@ -349,6 +395,7 @@ def tensorleap_integration_test():
                                         f'Integration test is only allowed to call Tensorleap decorators. '
                                         f'Ensure any arithmetics, external library use, Python logic is placed within Tensorleap decoders')
                 finally:
+                    _mapping_dataset_is_grouped = False
                     if mapping_runtime_mode_env_var_mame in os.environ:
                         del os.environ[mapping_runtime_mode_env_var_mame]
             finally:
@@ -639,10 +686,29 @@ def tensorleap_load_model(prediction_types: Optional[List[PredictionTypeHandler]
 
             class ModelPlaceholder:
 
+                @staticmethod
+                def _reject_reused_input_source(elem, seen):
+                    # A model input slot is wired by mutating elem.node_mapping.type, so the same
+                    # source object at two positions silently overwrites the first slot. Each input
+                    # encoder mints its own NodeMapping, so a collision means one source fed two
+                    # slots — e.g. indexing a single grouped input, model([imgs[0], imgs[1]]), where
+                    # imgs[i] returns the same node (a group is the batch dim of ONE input, not many
+                    # inputs), or the flat model([img, img]). Reject it instead of mis-wiring.
+                    if id(elem.node_mapping) in seen:
+                        raise Exception(
+                            "The same input source was passed to more than one model input slot. Each "
+                            "model input must come from a distinct input encoder. Indexing a single "
+                            "grouped input (e.g. model([imgs[0], imgs[1]])) feeds several samples of ONE "
+                            "input into separate slots — a group is the batch dimension of one input, "
+                            "not multiple inputs.")
+                    seen.add(id(elem.node_mapping))
+
                 # keras interface
                 def __call__(self, arg):
                     if isinstance(arg, list):
+                        seen: set = set()
                         for i, elem in enumerate(arg):
+                            self._reject_reused_input_source(elem, seen)
                             elem.node_mapping.type = _safe_get_item(i)
                     else:
                         arg.node_mapping.type = NodeMappingType.Input0
@@ -656,7 +722,9 @@ def tensorleap_load_model(prediction_types: Optional[List[PredictionTypeHandler]
                                                   "use model.run(None, inputs) to return all outputs.")
                     assert isinstance(input_dict, dict), \
                         f'Expected input_dict to be a dict, got {type(input_dict)} instead.'
+                    seen: set = set()
                     for i, (input_key, elem) in enumerate(input_dict.items()):
+                        self._reject_reused_input_source(elem, seen)
                         if isinstance(input_key, NodeMappingType):
                             elem.node_mapping.type = input_key
                         else:
@@ -1561,13 +1629,10 @@ def tensorleap_metadata(
                 raise Exception(f'Metadata with name {name} already exists. '
                                 f'Please choose another')
 
-        def _validate_input_args(sample_id: Union[int, str], preprocess_response: PreprocessResponse):
-            assert type(sample_id) == preprocess_response.sample_id_type, \
-                (f'{user_function.__name__}() validation failed: '
-                 f'Argument sample_id should be as the same type as defined in the preprocess response '
-                 f'{preprocess_response.sample_id_type}. Got {type(sample_id)}.')
+        def _validate_input_args(sample_id: Union[int, str, list], preprocess_response: PreprocessResponse):
+            _validate_id_or_group(sample_id, preprocess_response, user_function.__name__)
 
-        def _validate_result(result):
+        def _validate_single(result):
             supported_result_types = (type(None), int, str, bool, float, dict, np.floating,
                                       np.bool_, np.unsignedinteger, np.signedinteger, np.integer)
             if isinstance(result, tuple):
@@ -1585,6 +1650,20 @@ def tensorleap_metadata(
                         (f'{user_function.__name__}() validation failed: '
                          f'Values in the return dict should be of type {str(supported_result_types)}. Got {type(value)}.')
 
+        def _validate_result(result, grouped=False, group_size=None):
+            if not grouped:
+                _validate_single(result)
+                return
+            # Grouped metadata: a list of `group_size` per-sample scalars/dicts.
+            assert isinstance(result, list), \
+                (f'{user_function.__name__}() validation failed: expected a grouped output for a group '
+                 f'of {group_size}: a list of {group_size} metadata values, got {type(result)}.')
+            assert len(result) == group_size, \
+                (f'{user_function.__name__}() validation failed: expected a grouped output for a group '
+                 f'of {group_size}: a list of {group_size} metadata values, got {len(result)}.')
+            for value in result:
+                _validate_single(value)
+
         def inner_without_validate(sample_id, preprocess_response):
 
             global _called_from_inside_tl_decorator
@@ -1604,14 +1683,16 @@ def tensorleap_metadata(
                 set_current('tensorleap_metadata')
             if os.environ.get(mapping_runtime_mode_env_var_mame):
                 return None
-            validate_args_structure(*args, types_order=[Union[int, str], PreprocessResponse],
+            validate_args_structure(*args, types_order=[Union[int, str, list], PreprocessResponse],
                                     func_name=user_function.__name__, expected_names=["idx", "preprocess"], **kwargs)
             sample_id, preprocess_response = args if len(args) != 0 else kwargs.values()
             _validate_input_args(sample_id, preprocess_response)
 
+            grouped = isinstance(sample_id, list)
+            group_size = len(sample_id) if grouped else None
             result = inner_without_validate(sample_id, preprocess_response)
 
-            _validate_result(result)
+            _validate_result(result, grouped, group_size)
             if not _call_from_tl_platform:
                 update_env_params_func("tensorleap_metadata", "v")
             return result
@@ -1621,38 +1702,58 @@ def tensorleap_metadata(
     return decorating_function
 
 
-def tensorleap_custom_latent_space(name: Optional[str] = None):
+def tensorleap_custom_latent_space(name: Optional[str] = None, use_ls_for_analysis: bool = False):
+    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)}.")
+
     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: '
-                 f'Argument sample_id should be either int or str. Got {type(sample_id)}.')
-            assert isinstance(preprocess_response, PreprocessResponse), \
-                (f'tensorleap_custom_latent_space validation failed: '
-                 f'Argument preprocess_response should be a PreprocessResponse. Got {type(preprocess_response)}.')
-            assert type(sample_id) == preprocess_response.sample_id_type, \
-                (f'tensorleap_custom_latent_space validation failed: '
-                 f'Argument sample_id should be as the same type as defined in the preprocess response '
-                 f'{preprocess_response.sample_id_type}. Got {type(sample_id)}.')
 
-        def _validate_result(result):
-            assert isinstance(result, np.ndarray), \
+        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')
+
+        def _validate_single(single_result):
+            assert isinstance(single_result, np.ndarray), \
                 (f'tensorleap_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))
+                 f'The return type should be a numpy array. Got {type(single_result)}.')
+            if single_result.ndim > 1:
+                flat_dim = int(np.prod(single_result.shape))
                 store_general_warning(
-                    key=("tensorleap_custom_latent_space_flatten", ls_name, tuple(result.shape)),
+                    key=("tensorleap_custom_latent_space_flatten", ls_name, tuple(single_result.shape)),
                     message=(
-                        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"tensorleap_custom_latent_space '{ls_name}' returned per-sample shape {tuple(single_result.shape)} "
+                        f"(ndim={single_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 "
                         f"inside your function."
                     ),
                 )
 
+        def _validate_result(result, grouped=False, group_size=None):
+            if not grouped:
+                _validate_single(result)
+                return
+            # Grouped: (B, *) array or list of B per-sample arrays. Validate each per-sample slice so
+            # the per-sample flatten check/warning uses the sample shape, not the (B, *) batch shape.
+            if isinstance(result, np.ndarray):
+                assert result.ndim >= 1 and result.shape[0] == group_size, \
+                    (f'tensorleap_custom_latent_space validation failed: expected a grouped output for a '
+                     f'group of {group_size}: an array with leading dim {group_size}, got shape {result.shape}.')
+                for single in result:
+                    _validate_single(single)
+            elif isinstance(result, list):
+                assert len(result) == group_size, \
+                    (f'tensorleap_custom_latent_space validation failed: expected a grouped output for a '
+                     f'group of {group_size}: a list of {group_size} arrays, got {len(result)}.')
+                for single in result:
+                    _validate_single(single)
+            else:
+                raise AssertionError(
+                    f'tensorleap_custom_latent_space validation failed: expected a grouped output for a '
+                    f'group of {group_size}: either an array with leading dim {group_size} or a list of '
+                    f'{group_size} arrays, got {type(result)}.')
+
         def inner_without_validate(sample_id, preprocess_response):
             global _called_from_inside_tl_decorator
             _called_from_inside_tl_decorator += 1
@@ -1664,7 +1765,8 @@ def tensorleap_custom_latent_space(name: Optional[str] = None):
 
             return result
 
-        leap_binder.set_custom_latent_space(inner_without_validate, ls_name)
+        leap_binder.set_custom_latent_space(inner_without_validate, ls_name,
+                                            use_ls_for_analysis=use_ls_for_analysis)
 
         def inner(sample_id, preprocess_response):
             if os.environ.get(mapping_runtime_mode_env_var_mame):
@@ -1672,9 +1774,11 @@ def tensorleap_custom_latent_space(name: Optional[str] = None):
 
             _validate_input_args(sample_id, preprocess_response)
 
+            grouped = isinstance(sample_id, list)
+            group_size = len(sample_id) if grouped else None
             result = inner_without_validate(sample_id, preprocess_response)
 
-            _validate_result(result)
+            _validate_result(result, grouped, group_size)
             return result
 
         return inner
@@ -2643,6 +2747,15 @@ def tensorleap_preprocess():
             assert len(set(result)) == len(result), \
                 (f'{user_function.__name__}() validation failed: '
                  f'The return list should not contain duplicate PreprocessResponse objects.')
+            # All-or-none: a dataset is either all grouped (nested sample_ids) or all flat.
+            if len({response.is_grouped for response in result}) != 1:
+                shapes = ', '.join(
+                    f"#{i}={'grouped' if response.is_grouped else 'flat'}"
+                    for i, response in enumerate(result))
+                raise AssertionError(
+                    f'{user_function.__name__}() validation failed: '
+                    f'All PreprocessResponses must be either all grouped (nested sample_ids) '
+                    f'or all flat, but got a mix: {shapes}.')
 
         def inner(*args, **kwargs):
             if not _call_from_tl_platform:
@@ -2760,6 +2873,13 @@ def tensorleap_element_instance_preprocess(
             result = user_function()
             found_instance_metadata = False
             for preprocess_response in result:
+                if preprocess_response.is_grouped:
+                    raise Exception(
+                        "tensorleap_element_instance_preprocess validation failed: a grouped "
+                        "(nested sample_ids) PreprocessResponse cannot also declare element "
+                        "instances -- instance ids are derived per scalar sample id, and the "
+                        "engine's grouped/batch fetch path has no per-instance id to mask by. "
+                        "Remove the instance encoders or don't group this dataset.")
                 sample_ids_to_instance_mappings = {}
                 instance_to_sample_ids_mappings = {}
                 all_sample_ids = preprocess_response.sample_ids.copy()
@@ -2980,14 +3100,10 @@ def tensorleap_input_encoder(name: str, channel_dim=_UNSET, model_input_index=No
                 f"channel axis in the batched tensor (batch = axis 0), so 0 is invalid. Use 1 "
                 f"for channel-first (NCHW) and -1 for channel-last (NHWC).")
 
-        def _validate_input_args(sample_id: Union[int, str], preprocess_response: PreprocessResponse):
-            assert type(sample_id) == preprocess_response.sample_id_type, \
-                (f'{user_function.__name__}() validation failed: '
-                 f'Argument sample_id should be as the same type as defined in the preprocess response '
-                 f'{preprocess_response.sample_id_type}. Got {type(sample_id)}.')
+        def _validate_input_args(sample_id: Union[int, str, list], preprocess_response: PreprocessResponse):
+            _validate_id_or_group(sample_id, preprocess_response, user_function.__name__)
 
-        def _validate_result(result):
-            validate_output_structure(result, func_name=user_function.__name__, expected_type_name="np.ndarray")
+        def _validate_single(result):
             assert isinstance(result, np.ndarray), \
                 (f'{user_function.__name__}() validation failed: '
                  f'Unsupported return type. Should be a numpy array. Got {type(result)}.')
@@ -2997,6 +3113,13 @@ def tensorleap_input_encoder(name: str, channel_dim=_UNSET, model_input_index=No
             assert channel_dim - 1 <= len(result.shape), (f'{user_function.__name__}() validation failed: '
                                                           f'The channel_dim ({channel_dim}) should be <= to the rank of the resulting input rank ({len(result.shape)}).')
 
+        def _validate_result(result, grouped=False, group_size=None):
+            if not grouped:
+                validate_output_structure(result, func_name=user_function.__name__, expected_type_name="np.ndarray")
+                _validate_single(result)
+                return
+            _validate_grouped_result(result, group_size, user_function.__name__, _validate_single)
+
         def inner_without_validate(sample_id, preprocess_response):
             global _called_from_inside_tl_decorator
             _called_from_inside_tl_decorator += 1
@@ -3013,18 +3136,29 @@ def tensorleap_input_encoder(name: str, channel_dim=_UNSET, model_input_index=No
         def inner(*args, **kwargs):
             if not _call_from_tl_platform:
                 set_current("tensorleap_input_encoder")
-            validate_args_structure(*args, types_order=[Union[int, str], PreprocessResponse],
+            validate_args_structure(*args, types_order=[Union[int, str, list], PreprocessResponse],
                                     func_name=user_function.__name__, expected_names=["idx", "preprocess"], **kwargs)
             sample_id, preprocess_response = args if len(args) != 0 else kwargs.values()
             _validate_input_args(sample_id, preprocess_response)
 
+            grouped = isinstance(sample_id, list)
+            group_size = len(sample_id) if grouped else None
             result = inner_without_validate(sample_id, preprocess_response)
 
-            _validate_result(result)
+            _validate_result(result, grouped, group_size)
 
             if _called_from_inside_tl_decorator == 0 and _called_from_inside_tl_integration_test_decorator:
-                batch_warning(result, user_function.__name__)
-                result = np.expand_dims(result, axis=0)
+                if grouped:
+                    # A grouped result is a list of per-sample arrays (never stacked, so a
+                    # B=1-only model can be fed one sample at a time). A user may also return a
+                    # supported (B, *) array; split it back into per-sample rows first (mirrors
+                    # LeapLoader._to_grouped_list) so both forms match the runtime path. Then add
+                    # a batch dim to each sample, mirroring the flat path's single-sample expand_dims.
+                    rows = result if isinstance(result, list) else [row for row in np.asarray(result)]
+                    result = [np.expand_dims(r, axis=0) for r in rows]
+                else:
+                    batch_warning(result, user_function.__name__)
+                    result = np.expand_dims(result, axis=0)
                 # Emit integration test event once per test
                 try:
                     emit_integration_event_once(AnalyticsEvent.INPUT_ENCODER_INTEGRATION_TEST, {
@@ -3045,8 +3179,15 @@ def tensorleap_input_encoder(name: str, channel_dim=_UNSET, model_input_index=No
         inner.node_mapping = NodeMapping(name, node_mapping_type)
 
         def mapping_inner(*args, **kwargs):
-            class TempMapping:
-                pass
+            if _mapping_dataset_is_grouped:
+                class TempMapping:
+                    def __getitem__(self, key):
+                        # Indexing a grouped input group (input_list[i]) selects a sample from the
+                        # same input source; the model wiring is identical to the un-grouped input.
+                        return self
+            else:
+                class TempMapping:
+                    pass
 
             ret = TempMapping()
             ret.node_mapping = mapping_inner.node_mapping
@@ -3076,15 +3217,10 @@ def tensorleap_gt_encoder(name: str):
                 raise Exception(f'GT with name {name} already exists. '
                                 f'Please choose another')
 
-        def _validate_input_args(sample_id: Union[int, str], preprocess_response: PreprocessResponse):
-            assert type(sample_id) == preprocess_response.sample_id_type, \
-                (f'{user_function.__name__}() validation failed: '
-                 f'Argument sample_id should be as the same type as defined in the preprocess response '
-                 f'{preprocess_response.sample_id_type}. Got {type(sample_id)}.')
+        def _validate_input_args(sample_id: Union[int, str, list], preprocess_response: PreprocessResponse):
+            _validate_id_or_group(sample_id, preprocess_response, user_function.__name__)
 
-        def _validate_result(result):
-            validate_output_structure(result, func_name=user_function.__name__, expected_type_name="np.ndarray",
-                                      gt_flag=True)
+        def _validate_single(result):
             assert isinstance(result, np.ndarray), \
                 (f'{user_function.__name__}() validation failed: '
                  f'Unsupported return type. Should be a numpy array. Got {type(result)}.')
@@ -3092,6 +3228,14 @@ def tensorleap_gt_encoder(name: str):
                 (f'{user_function.__name__}() validation failed: '
                  f'The return type should be a numpy array of type float32. Got {result.dtype}.')
 
+        def _validate_result(result, grouped=False, group_size=None):
+            if not grouped:
+                validate_output_structure(result, func_name=user_function.__name__, expected_type_name="np.ndarray",
+                                          gt_flag=True)
+                _validate_single(result)
+                return
+            _validate_grouped_result(result, group_size, user_function.__name__, _validate_single)
+
         def inner_without_validate(sample_id, preprocess_response):
             global _called_from_inside_tl_decorator
             _called_from_inside_tl_decorator += 1
@@ -3108,18 +3252,29 @@ def tensorleap_gt_encoder(name: str):
         def inner(*args, **kwargs):
             if not _call_from_tl_platform:
                 set_current("tensorleap_gt_encoder")
-            validate_args_structure(*args, types_order=[Union[int, str], PreprocessResponse],
+            validate_args_structure(*args, types_order=[Union[int, str, list], PreprocessResponse],
                                     func_name=user_function.__name__, expected_names=["idx", "preprocess"], **kwargs)
-            sample_id, preprocess_response = args
+            sample_id, preprocess_response = args if len(args) != 0 else kwargs.values()
             _validate_input_args(sample_id, preprocess_response)
 
+            grouped = isinstance(sample_id, list)
+            group_size = len(sample_id) if grouped else None
             result = inner_without_validate(sample_id, preprocess_response)
 
-            _validate_result(result)
+            _validate_result(result, grouped, group_size)
 
             if _called_from_inside_tl_decorator == 0 and _called_from_inside_tl_integration_test_decorator:
-                batch_warning(result, user_function.__name__)
-                result = np.expand_dims(result, axis=0)
+                if grouped:
+                    # A grouped result is a list of per-sample arrays (never stacked, so a
+                    # B=1-only model can be fed one sample at a time). A user may also return a
+                    # supported (B, *) array; split it back into per-sample rows first (mirrors
+                    # LeapLoader._to_grouped_list) so both forms match the runtime path. Then add
+                    # a batch dim to each sample, mirroring the flat path's single-sample expand_dims.
+                    rows = result if isinstance(result, list) else [row for row in np.asarray(result)]
+                    result = [np.expand_dims(r, axis=0) for r in rows]
+                else:
+                    batch_warning(result, user_function.__name__)
+                    result = np.expand_dims(result, axis=0)
                 _register_chain_artifact(result, 'ground_truth')
                 # Emit integration test event once per test
                 try:
@@ -3135,8 +3290,15 @@ def tensorleap_gt_encoder(name: str):
         inner.node_mapping = NodeMapping(name, NodeMappingType.GroundTruth)
 
         def mapping_inner(*args, **kwargs):
-            class TempMapping:
-                pass
+            if _mapping_dataset_is_grouped:
+                class TempMapping:
+                    def __getitem__(self, key):
+                        # Indexing a grouped GT group (gt_list[i]) selects a sample from the same
+                        # ground-truth source; wiring is identical to the un-grouped GT node.
+                        return self
+            else:
+                class TempMapping:
+                    pass
 
             ret = TempMapping()
             ret.node_mapping = mapping_inner.node_mapping