fix: slice batch_prob per mask instead of indexing it

smauglab/transforms/gpu/base.py

@@ -326,7 +326,12 @@ def transform_list(
             params = cls.get_instance_module_param(param)
             params_i = copy.deepcopy(params)
             for i, inp in enumerate(input):
-                params_i["batch_prob"] = params["batch_prob"][i]
+                # [i : i + 1], not [i]: indexing with a scalar gives a 0-dim
+                # tensor, and kornia does `to_apply = batch_prob > 0.5` and then
+                # branches on `to_apply.all()` / `.any()`. Those are the same value
+                # for a 0-dim tensor, so the per-element branch was unreachable and
+                # `in_tensor[to_apply]` would select along the wrong axis.
+                params_i["batch_prob"] = params["batch_prob"][i : i + 1]
                 tfm_inp = module.transform_masks(inp, params=params_i, flags=module.flags, transform=module.transform_matrix, **extra_args)
                 tfm_input.append(tfm_inp)
             input = tfm_input
@@ -336,7 +341,12 @@ def transform_list(
             params = cls.get_instance_module_param(param)
             params_i = copy.deepcopy(params)
             for i, inp in enumerate(input):
-                params_i["batch_prob"] = params["batch_prob"][i]
+                # [i : i + 1], not [i]: indexing with a scalar gives a 0-dim
+                # tensor, and kornia does `to_apply = batch_prob > 0.5` and then
+                # branches on `to_apply.all()` / `.any()`. Those are the same value
+                # for a 0-dim tensor, so the per-element branch was unreachable and
+                # `in_tensor[to_apply]` would select along the wrong axis.
+                params_i["batch_prob"] = params["batch_prob"][i : i + 1]
                 tfm_inp = module.transform_masks(inp, params=params_i, flags=module.flags, **extra_args)
                 tfm_input.append(tfm_inp)
             input = tfm_input

unit_tests/test_mask_list_batch_prob.py

@@ -0,0 +1,58 @@
+"""`MaskSequentialOpsCustom.transform_list` must hand kornia a 1-element batch.
+
+The per-mask loop sliced the application mask with `params["batch_prob"][i]`,
+which gives a **0-dim** tensor. Kornia then does
+
+    to_apply = batch_prob > 0.5
+    if to_apply.all(): ... elif not to_apply.any(): ... else: ...
+
+and for a 0-dim tensor `.all()` and `.any()` are always equal, so the
+"some samples apply" branch is unreachable -- and `in_tensor[to_apply]` with a
+0-dim boolean indexes along the wrong axis instead of selecting rows.
+
+This path is latent today: `AugmentationSequentialOpsCustom.transform` asserts
+the mask is a Tensor before the list branch is reached, so nothing in SmaugLab
+can get here. The slice is still wrong, and the assert is the only thing
+standing between it and a caller who passes a list of masks -- which is the
+signature's whole reason to exist.
+"""
+
+from __future__ import annotations
+
+import inspect
+
+import torch
+
+from smauglab.transforms.gpu import base
+from unit_tests.helpers import SmaugLabTestCase
+
+
+class TestListMaskBatchProb(SmaugLabTestCase):
+    def test_the_slice_keeps_the_batch_axis(self):
+        """What kornia needs, and what the two forms actually give."""
+        probs = torch.tensor([1.0, 0.0, 1.0])
+
+        self.assertEqual(probs[0:1].dim(), 1, "a 1-element slice keeps the batch axis")
+        self.assertEqual(probs[0].dim(), 0, "a scalar index drops it")
+
+    def test_a_zero_dim_mask_collapses_kornias_branching(self):
+        """Why the 0-dim form is not merely untidy."""
+        scalar = torch.tensor([1.0, 0.0])[1] > 0.5
+
+        self.assertEqual(bool(scalar.all()), bool(scalar.any()), "all() and any() coincide, so the partial branch is dead")
+
+    def test_the_loop_slices_rather_than_indexes(self):
+        """Pinned on the source, because the path is unreachable from the pipeline.
+
+        `AugmentationSequentialOpsCustom.transform` asserts the mask is a Tensor,
+        so a list never reaches `transform_list`. Asserting on the text is worth
+        more than a test that cannot run the code.
+        """
+        source = inspect.getsource(base.MaskSequentialOpsCustom.transform_list)
+
+        self.assertNotIn('params["batch_prob"][i]', source, "the scalar index is back")
+        self.assertEqual(
+            source.count('params["batch_prob"][i : i + 1]'),
+            2,
+            "both per-mask loops must slice the application mask",
+        )