Repository: pystiche/pystiche

Issue:
Add target_image, target_guide, input_guide property to SameOperatorContainer
Consider the following setup:

```python
import torch
from torch import nn
from pystiche import enc, ops


encoder = enc.SequentialEncoder((nn.Conv2d(3, 3, 1),))
names = ("conv",)


def get_op(name, layer_weight):
    return ops.GramOperator(encoder, score_weight=layer_weight)


op = ops.SameOperatorContainer(names, get_op)

target_image = torch.rand(1, 3, 32, 32)
```

I can conveniently `set_target_image` (as well as `set_target_guide` and `set_input_guide`)

```python
op.set_target_image(target_image)
```

but to retrieve it I have to know the internal structure of `op` (here knowing that `op` has a sub-operator named `conv`)

```python
retrieved_target_image = op.conv.target_image
assert retrieved_target_image is target_image
```

---

I propose a `target_image` (as well as `target_guide` and `target_image`) property that retrieves the image. If there are multiple sub-operators, as is usually the case, we can collect all images, check if they are the same and return the image. If they don't match we could raise an error stating just this.


Discussion:

Resolving patch:
diff --git a/pystiche/ops/container.py b/pystiche/ops/container.py
--- a/pystiche/ops/container.py
+++ b/pystiche/ops/container.py
@@ -1,5 +1,5 @@
 from collections import OrderedDict
-from typing import Callable, Dict, Sequence, Tuple, Union, cast
+from typing import Any, Callable, Dict, List, Sequence, Tuple, Union, cast
 
 import torch
 from torch import nn
@@ -40,18 +40,92 @@ def process_input_image(self, input_image: torch.Tensor) -> pystiche.LossDict:
             [(name, op(input_image)) for name, op in self.named_children()]
         )
 
+    def _get_image_or_guide(
+        self, attr: str, comparison_only: bool = False
+    ) -> torch.Tensor:
+        images_or_guides: List[torch.Tensor] = []
+        for op in self.operators():
+            if comparison_only and not isinstance(op, ComparisonOperator):
+                continue
+
+            try:
+                images_or_guides.append(getattr(op, attr))
+            except AttributeError:
+                pass
+
+        if not images_or_guides:
+            raise RuntimeError(f"No immediate children has a {attr}.")
+
+        image_or_guide = images_or_guides[0]
+        key = pystiche.TensorKey(image_or_guide)
+
+        if not all(key == other for other in images_or_guides[1:]):
+            raise RuntimeError(f"The immediate children have non-matching {attr}")
+
+        return image_or_guide
+
+    def get_target_guide(self) -> torch.Tensor:
+        r"""Extracts the target guide from the immediate children
+
+        Returns:
+            guide: Target guide of shape :math:`1 \times 1 \times H \times W`.
+
+        Raises:
+            RuntimeError: If no immediate children has a target guide or the the target
+                guides do not match each other
+        """
+        return self._get_image_or_guide("target_guide", comparison_only=True)
+
+    def get_target_image(self) -> torch.Tensor:
+        r"""Extracts the target image from the immediate children
+
+        Returns:
+            image: Target image of shape :math:`1 \times 1 \times H \times W`.
+
+        Raises:
+            RuntimeError: If no immediate children has a target image or the the target
+                images do not match each other
+        """
+        return self._get_image_or_guide("target_image", comparison_only=True)
+
+    def get_input_guide(self) -> torch.Tensor:
+        r"""Extracts the input guide from the immediate children
+
+        Returns:
+            guide: Input guide of shape :math:`1 \times 1 \times H \times W`.
+
+        Raises:
+            RuntimeError: If no immediate children has a input guide or the the input
+                guides do not match each other
+        """
+        return self._get_image_or_guide("input_guide")
+
+    def _set_image_or_guide(
+        self,
+        image_or_guide: torch.Tensor,
+        attr: str,
+        comparison_only: bool = False,
+        **kwargs: Any,
+    ) -> None:
+        for op in self.operators():
+            if comparison_only and not isinstance(op, ComparisonOperator):
+                continue
+
+            setter = getattr(op, f"set_{attr}")
+            setter(image_or_guide, **kwargs)
+
     def set_target_guide(self, guide: torch.Tensor, recalc_repr: bool = True) -> None:
         r"""Invoke :meth:`~pystiche.ops.ComparisonOperator.set_target_guide` on all
         immediate :class:`~pystiche.ops.ComparisonOperator` children.
 
         Args:
-            guide: Input guide of shape :math:`1 \times 1 \times H \times W`.
+            guide: Target guide of shape :math:`1 \times 1 \times H \times W`.
             recalc_repr: If ``True``, recalculates :meth:`.target_enc_to_repr`.
                 Defaults to ``True``.
         """
-        for op in self.operators():
-            if isinstance(op, ComparisonOperator):
-                op.set_target_guide(guide, recalc_repr=recalc_repr)
+        self._set_image_or_guide(
+            guide, "target_guide", comparison_only=True, recalc_repr=recalc_repr
+        )
 
     def set_target_image(self, image: torch.Tensor) -> None:
         r"""Invoke :meth:`~pystiche.ops.ComparisonOperator.set_target_image` on all
@@ -60,9 +134,7 @@ def set_target_image(self, image: torch.Tensor) -> None:
         Args:
             image: Target image of shape :math:`B \times C \times H \times W`.
         """
-        for op in self.operators():
-            if isinstance(op, ComparisonOperator):
-                op.set_target_image(image)
+        self._set_image_or_guide(image, "target_image", comparison_only=True)
 
     def set_input_guide(self, guide: torch.Tensor) -> None:
         r"""Invoke :meth:`~pystiche.ops.Operator.set_input_guide` on all immediate
@@ -71,8 +143,7 @@ def set_input_guide(self, guide: torch.Tensor) -> None:
         Args:
             guide: Input guide of shape :math:`1 \times 1 \times H \times W`.
         """
-        for op in self.operators():
-            op.set_input_guide(guide)
+        self._set_image_or_guide(guide, "input_guide")
 
     _modules: Dict[str, nn.Module]
 


Test patch:
diff --git a/tests/test_ops.py b/tests/test_ops.py
--- a/tests/test_ops.py
+++ b/tests/test_ops.py
@@ -219,6 +219,74 @@ def process_input_image(self, image):
         desireds = named_ops
         self.assertNamedChildrenEqual(actuals, desireds)
 
+    def test_OperatorContainer_get_image_or_guide(self):
+        class RegularizationTestOperator(ops.PixelRegularizationOperator):
+            def input_image_to_repr(self, image):
+                pass
+
+            def calculate_score(self, input_repr):
+                pass
+
+        class ComparisonTestOperator(ops.PixelComparisonOperator):
+            def target_image_to_repr(self, image):
+                return image, None
+
+            def input_image_to_repr(self, image, ctx):
+                pass
+
+            def calculate_score(self, input_repr, target_repr, ctx):
+                pass
+
+        def get_container():
+            return ops.OperatorContainer(
+                (
+                    ("regularization", RegularizationTestOperator()),
+                    ("comparison", ComparisonTestOperator()),
+                )
+            )
+
+        torch.manual_seed(0)
+        image_or_guide = torch.rand(1, 3, 128, 128)
+
+        with self.subTest("no image or guide"):
+            container = get_container()
+
+            with self.assertRaises(RuntimeError):
+                container.get_input_guide()
+
+        with self.subTest("mismatching images or guides"):
+            container = get_container()
+            container.regularization.set_input_guide(image_or_guide)
+            container.comparison.set_input_guide(-image_or_guide)
+
+            with self.assertRaises(RuntimeError):
+                container.get_input_guide()
+
+        with self.subTest("get_target_guide"):
+            container = get_container()
+            container.comparison.set_target_guide(image_or_guide)
+
+            actual = container.get_target_guide()
+            desired = image_or_guide
+            self.assertTensorAlmostEqual(actual, desired)
+
+        with self.subTest("get_target_image"):
+            container = get_container()
+            container.comparison.set_target_image(image_or_guide)
+
+            actual = container.get_target_image()
+            desired = image_or_guide
+            self.assertTensorAlmostEqual(actual, desired)
+
+        with self.subTest("get_input_guide"):
+            container = get_container()
+            container.regularization.set_input_guide(image_or_guide)
+            container.comparison.set_input_guide(image_or_guide)
+
+            actual = container.get_input_guide()
+            desired = image_or_guide
+            self.assertTensorAlmostEqual(actual, desired)
+
     def test_OperatorContainer_set_image_or_guide(self):
         class RegularizationTestOperator(ops.PixelRegularizationOperator):
             def input_image_to_repr(self, image):