Better type handling for rshifts (#369)

**Context:**

```
from mrmustard.lab_dev import Coherent

Coherent([0, 1]) >> Coherent([0]).dual
```
At the moment, the code above returns a `CircuitComponent`. We want it
to return a `Ket`.

mrmustard/lab_dev/states/base.py

@@ -34,7 +34,6 @@
 from mrmustard.physics.gaussian import purity
 from mrmustard.physics.representations import Bargmann, Fock
 from ..circuit_components import CircuitComponent
-from ..transformations.transformations import Unitary, Channel
 
 __all__ = ["State", "DM", "Ket"]
 
@@ -351,17 +350,14 @@ def __rshift__(self, other: CircuitComponent) -> CircuitComponent:
         Contracts ``self`` and ``other`` as it would in a circuit, adding the adjoints when
         they are missing.
 
-        Returns a ``DM`` when ``other`` is a ``Unitary`` or a ``Channel``, and ``other`` acts on
-        ``self``'s modes. Otherwise, it returns a ``CircuitComponent``.
+        Returns a ``DM`` when the wires of the resulting components are compatible with those
+        of a ``Ket``, a ``CircuitComponent`` otherwise.
         """
-        component = super().__rshift__(other)
+        ret = super().__rshift__(other)
 
-        if isinstance(other, (Unitary, Channel)) and set(other.modes).issubset(self.modes):
-            dm = DM()
-            dm._wires = component.wires
-            dm._representation = component.representation
-            return dm
-        return component
+        if not ret.wires.input and ret.wires.bra.modes == ret.wires.ket.modes:
+            return DM._from_attributes("", ret.representation, ret.wires)
+        return ret
 
     def __repr__(self) -> str:
         return super().__repr__().replace("CircuitComponent", "DM")
@@ -478,23 +474,17 @@ def __rshift__(self, other: CircuitComponent) -> CircuitComponent:
         Contracts ``self`` and ``other`` as it would in a circuit, adding the adjoints when
         they are missing.
 
-        Returns a ``State`` (either ``Ket`` or ``DM``) when ``other`` is a ``Unitary`` or a
-        ``Channel``, and ``other`` acts on ``self``'s modes. Otherwise, it returns a
-        ``CircuitComponent``.
+        Returns a ``DM`` or a ``Ket`` when the wires of the resulting components are compatible
+        with those of a ``DM`` or of a ``Ket``, a ``CircuitComponent`` otherwise.
         """
-        component = super().__rshift__(other)
-
-        if isinstance(other, Unitary) and set(other.modes).issubset(set(self.modes)):
-            ket = Ket()
-            ket._wires = component.wires
-            ket._representation = component.representation
-            return ket
-        elif isinstance(other, Channel) and set(other.modes).issubset(set(self.modes)):
-            dm = DM()
-            dm._wires = component.wires
-            dm._representation = component.representation
-            return dm
-        return component
+        ret = super().__rshift__(other)
+
+        if not ret.wires.input:
+            if not ret.wires.bra:
+                return Ket._from_attributes("", ret.representation, ret.wires)
+            if ret.wires.bra.modes == ret.wires.ket.modes:
+                return DM._from_attributes("", ret.representation, ret.wires)
+        return ret
 
     def __repr__(self) -> str:
         return super().__repr__().replace("CircuitComponent", "Ket")

mrmustard/lab_dev/transformations/base.py

@@ -59,19 +59,13 @@ def __rshift__(self, other: CircuitComponent) -> CircuitComponent:
         Returns a ``Unitary`` when ``other`` is a ``Unitary``, a ``Channel`` when ``other`` is a
         ``Channel``, and a ``CircuitComponent`` otherwise.
         """
-        component = super().__rshift__(other)
+        ret = super().__rshift__(other)
 
         if isinstance(other, Unitary):
-            unitary = Unitary()
-            unitary._wires = component.wires
-            unitary._representation = component.representation
-            return unitary
+            return Unitary._from_attributes("", ret.representation, ret.wires)
         elif isinstance(other, Channel):
-            channel = Channel()
-            channel._wires = component.wires
-            channel._representation = component.representation
-            return channel
-        return component
+            return Channel._from_attributes("", ret.representation, ret.wires)
+        return ret
 
     def __repr__(self) -> str:
         return super().__repr__().replace("CircuitComponent", "Unitary")
@@ -101,14 +95,11 @@ def __rshift__(self, other: CircuitComponent) -> CircuitComponent:
         Returns a ``Channel`` when ``other`` is a ``Unitary`` or a ``Channel``, and a
         ``CircuitComponent`` otherwise.
         """
-        component = super().__rshift__(other)
+        ret = super().__rshift__(other)
 
         if isinstance(other, (Unitary, Channel)):
-            channel = Channel()
-            channel._wires = component.wires
-            channel._representation = component.representation
-            return channel
-        return component
+            return Channel._from_attributes("", ret.representation, ret.wires)
+        return ret
 
     def __repr__(self) -> str:
         return super().__repr__().replace("CircuitComponent", "Channel")

tests/test_lab_dev/test_states.py

@@ -152,13 +152,18 @@ def test_rshift(self):
             channel.name, channel.representation, channel.wires
         )  # pylint: disable=protected-access
 
+        # gates
         assert isinstance(ket >> unitary, Ket)
         assert isinstance(ket >> channel, DM)
         assert isinstance(ket >> unitary >> channel, DM)
         assert isinstance(ket >> channel >> unitary, DM)
         assert isinstance(ket >> u_component, CircuitComponent)
         assert isinstance(ket >> ch_component, CircuitComponent)
 
+        # measurements
+        assert isinstance(ket >> Coherent([0], 1).dual, Ket)
+        assert isinstance(ket >> Coherent([0], 1).dm().dual, DM)
+
     def test_repr(self):
         ket = Coherent([0, 1], 1)
         ket_component = CircuitComponent._from_attributes(
@@ -273,13 +278,18 @@ def test_rshift(self):
         )  # pylint: disable=protected-access
 
         dm = ket >> channel
-        assert isinstance(dm, DM)
 
+        # gates
+        assert isinstance(dm, DM)
         assert isinstance(dm >> unitary >> channel, DM)
         assert isinstance(dm >> channel >> unitary, DM)
         assert isinstance(dm >> u_component, CircuitComponent)
         assert isinstance(dm >> ch_component, CircuitComponent)
 
+        # measurements
+        assert isinstance(dm >> Coherent([0], 1).dual, DM)
+        assert isinstance(dm >> Coherent([0], 1).dm().dual, DM)
+
     def test_repr(self):
         ket = Coherent([0, 1], 1)
         channel = Attenuator([1], 1)