Repository: lf1-io/padl (language: python)

Issue:
Grouping not kept in `pd_preprocess`, `pd_forward` and `pd_postprocess`
When we have a Transform `g` compounded by several elements and any of them is grouped, this grouping is not respected in `g.pd_preprocess` `g.pd_forward`, `g.pd_postprocess`. For example,

Let `t`  and `g` be:
```
linear = transform(torch.nn.Linear)(8, 4)
t = ((to_tensor >> padl.Batchify() >>  linear)
     + (to_tensor >> padl.Batchify() >> linear)
) - 'name' 
g = t + padl.Identity()
```
`t` is:
``` 
Rollout - "name":

   ├─▶ 0: to_tensor >> Batchify(dim=0) >> <anonymous Linear>
   │ +  
   └─▶ 1: to_tensor >> Batchify(dim=0) >> <anonymous Linear>
```
`g` is:
```
Rollout:

   ├─▶ 0: group([..>>..] + [..>>..])
   │ +  
   └─▶ 1: padl.Identity()
```
And `g.pd_preprocess` and `g.pd_forward` are, respectively:
```
Rollout:

   ├─▶ 0: to_tensor >> Batchify(dim=0)
   │ +  
   ├─▶ 1: to_tensor >> Batchify(dim=0)
   │ +  
   └─▶ 2: padl.Identity()
```
,
```
Parallel:

   ││└─▶ 0: Linear(in_features=8, out_features=4)
   ││  /  
   │└──▶ 1: Linear(in_features=8, out_features=4)
   │   /  
   └───▶ 2: padl.Identity()
```
@sjrl and me have been checking this. Maybe it would be more consistent if the grouping were kept on the `pd_forward` and `pd_preprocess`, so they would be, respectively, something like
```
Rollout:

   ├─▶ 0: group(...)
   │ +  
   ├─▶ 1: padl.Identity()
```
and
```
Parallel:

   │└─▶ 0: group(...)
   │  /  
   └──▶ 1: padl.Identity()
```

Update: this also happens in Rollouts of Parallels and viceversa, for example in
```
t = ((to_tensor >> batch >> linear) / (to_tensor >> batch >> linear)
     + (to_tensor >> batch) / (to_tensor >> batch)
)
```
.
where `t.pd_preprocess` and `t._pd_forward` are, respectively,
```
Rollout:

   ├─▶ 0: [..>>..] / [..>>..]
   │ +  
   └─▶ 1: [..>>..] / [..>>..]
```
and 
```
Parallel:

   ││└─▶ 0: Linear(in_features=8, out_features=4)
   ││  /  
   │└──▶ 1: Linear(in_features=8, out_features=4)
   │   /  
   └───▶ 2: padl.Identity()
```
.



Discussion / hints:
Agree, I think it can be fixed by making grouped Pipelines automatically group their preprocess forward and postprocess.  
I reopen this issue because in our fix I think we missed the case for the Composes. When we have a `Pipeline` `g` which includes a `Compose` whose `pd_group` is `True` (mainly the case for this is because it has a name), this grouping is not respected in `g.pd_preprocess`, `g.pd_forward` and `g.pd_postprocess`.

For example, if we define `forward_step` as 
```
forward_step = (linear >> linear) - 'forward_step_1' 
```
`forward_step` is
```
Compose - "forward_step_1":

      │
      ▼ input
   0: Linear(in_features=8, out_features=4)
      │
      ▼ input
   1: Linear(in_features=8, out_features=4)
```

And let `g` be:
```
g = to_tensor >> batch >> forward_step >> linear 
```
Then,  `g` is:
```
Compose - "g":

      │
      ▼ x
   0: to_tensor                            
      │
      ▼ args
   1: Batchify(dim=0)                      
      │
      ▼ args
   2: Linear(in_features=8, out_features=4) >> Linear(in_features=8, out_features=4)
      │
      ▼ input
   3: Linear(in_features=8, out_features=4)

```
And we get that `g.pd_forward` is:
```
Compose:

      │
      ▼ input
   0: Linear(in_features=8, out_features=4)
      │
      ▼ input
   1: Linear(in_features=8, out_features=4)
      │
      ▼ input
   2: Linear(in_features=8, out_features=4)
```

Maybe it would make more sense and would be consistent to the previous fix if we kept the grouping of `forward_step` on `g.pd_forward`, so `g.pd_forward` would be something like
```
Compose:

      │
      ▼ input
   0: Linear(in_features=8, out_features=4) >> Linear(in_features=8, out_features=4)
      │
      ▼ input
   1: Linear(in_features=8, out_features=4)
```

Fix patch:
diff --git a/padl/transforms.py b/padl/transforms.py
index 8d39b4f2..dae616b5 100644
--- a/padl/transforms.py
+++ b/padl/transforms.py
@@ -1552,6 +1552,13 @@ def _pd_get_keys(transforms):
             deduped_keys.append(new_name)
         return deduped_keys
 
+    def _add_name_to_splits(self, final_splits):
+        """Ad name to split-transforms. """
+        if self._pd_name is not None:
+            for i, s in enumerate(final_splits):
+                if not isinstance(s, Identity):
+                    final_splits[i] = s - self._pd_name
+
 
 class Compose(Pipeline):
     """Apply series of transforms on input.
@@ -1623,6 +1630,8 @@ def _pd_splits(self, input_components=0, has_batchify=False) -> Tuple[Union[int,
             else:  # if it's empty: identity
                 final_splits.append(identity)
 
+        self._add_name_to_splits(final_splits)
+
         return output_components, final_splits, has_batchify
 
     @staticmethod
@@ -1837,8 +1846,6 @@ def _pd_splits(self, input_components=0) -> Tuple[Union[int, List],
             else:
                 final_splits.append(s)
 
-        final_splits = tuple(final_splits)
-
         res = []
         for split in final_splits:
             try:
@@ -1847,6 +1854,8 @@ def _pd_splits(self, input_components=0) -> Tuple[Union[int, List],
                 res.append(split)
         final_splits = res
 
+        self._add_name_to_splits(final_splits)
+
         return output_components, final_splits, has_batchify
 
     def __call__(self, args):
@@ -1940,6 +1949,8 @@ def _pd_splits(self, input_components=0) -> Tuple[Union[int, List],
                 res.append(split)
         final_splits = res
 
+        self._add_name_to_splits(final_splits)
+
         return output_components, final_splits, has_batchify
 
     def __call__(self, args):


Test patch:
diff --git a/tests/unittests/test_transforms.py b/tests/unittests/test_transforms.py
index f5ec2c67..60458b8c 100644
--- a/tests/unittests/test_transforms.py
+++ b/tests/unittests/test_transforms.py
@@ -713,9 +713,9 @@ def test_pd_splits_parallel(self, to_tensor, lin):
         g = t / Identity()
         t_preprocess = group(group((to_tensor >> batch) + (to_tensor >> batch)) /
                              group((to_tensor >> batch) + (to_tensor >> batch)))
-        g_preprocess = group(t_preprocess / Identity())
+        g_preprocess = group((t_preprocess - 'name') / Identity())
         t_forward = group(group(lin / lin) / Identity())
-        g_forward = group(t_forward / Identity())
+        g_forward = group((t_forward - 'name') / Identity())
         assert str(t.pd_preprocess) == str(t_preprocess)
         assert str(g.pd_preprocess) == str(g_preprocess)
         assert str(t.pd_forward) == str(t_forward)
@@ -726,9 +726,9 @@ def test_pd_splits_rollout(self, to_tensor, lin):
              + (to_tensor >> batch) / (to_tensor >> batch)
         ) - 'name'
         g = t + Identity()
-        t_preprocess = group(group(to_tensor / to_tensor) + group(to_tensor / to_tensor))
+        t_preprocess = group(group(to_tensor / to_tensor) + group(to_tensor / to_tensor)) - 'name'
         g_preprocess = group(t_preprocess + Identity())
-        t_forward = group(group(lin / lin) / Identity())
+        t_forward = group(group(lin / lin) / Identity()) - 'name'
         g_forward = group(t_forward / Identity())
         assert str(t.pd_preprocess) == str(t_preprocess)
         assert str(g.pd_preprocess) == str(g_preprocess)