Repository: CQCL/lambeq

Issue:
Tensor ansatze work only on `Word` boxes
Tensor ansatze (Spider, MPS) can be applied only on `Word` boxes, i.e. on boxes with empty domain. It would be desirable to make them work on arbitrary boxes, although this is not exactly trivial (input wires need to be split consistently between the resulting boxes).

Hints from issue comments:
@kinianlo: Would simply (un)currying be acceptable? Yes, cups are introduced but it shouldn't matter for tensor networks.

```python
from lambeq.backend.grammar import Box, Ty, Id, Diagram
from lambeq.backend.tensor import Dim
from lambeq.ansatz.tensor import MPSAnsatz, SpiderAnsatz

a, b, c, d, e = map(Ty, 'abcde')
box = Box('A', a @ b, c @ d @ e)
box.draw()
box_state = Box('A_state', Ty(), box.dom.r @ box.cod)
diag = Id(box.dom) @ box_state >> Diagram.cups(box.dom, box.dom.r) @ Id(box.cod) 
diag.draw()

mps_ansatz = MPSAnsatz({a: Dim(2), b: Dim(2), c: Dim(2), d: Dim(2), e: Dim(2)}, 10, 3)
spider_ansatz = SpiderAnsatz({a: Dim(2), b: Dim(2), c: Dim(2), d: Dim(2), e: Dim(2)}, 3)
mps_ansatz(diag).draw()
spider_ansatz(diag).draw()
```

![Image](https://github.com/user-attachments/assets/b1fe010a-1e24-455c-9b57-5e77f0ae7a23)

![Image](https://github.com/user-attachments/assets/7b7ce800-0ad5-4de4-9955-aa61b5983b6a)

![Image](https://github.com/user-attachments/assets/efa4b39c-0ab6-43d6-bd79-310c388166a3)

![Image](https://github.com/user-attachments/assets/a1db9685-ceca-44cf-aece-863f2b3d5e7a)@dimkart: Hi @kinianlo, thanks for this. Looks acceptable, however placing all inputs at the first box comes with certain risks, e.g. if any connector wire is discarded later or smth. But if done carefully could work I guess.@dimkart: I guess with this solution we'll require a kind of rewriter to do the uncurrying.@kinianlo: > Hi [@kinianlo](https://github.com/kinianlo), thanks for this. Looks acceptable, however placing all inputs at the first box comes with certain risks, e.g. if any connector wire is discarded later or smth. But if done carefully could work I guess.

Hi @dimkart , I'm not sure I understand the project well enough to spot these risks. It would be nice if there is an example.


> I guess with this solution we'll require a kind of rewriter to do the uncurrying.

Something like this? 
```python
class UncurryRewriteRule(RewriteRule):
    def __init__(self, left: bool) -> None:
        self.left = left # True if you want the cups on the left

    def matches(self, box: Box) -> bool:
        return bool(box.dom)

    def rewrite(self, box: Box) -> Diagrammable:
        if self.left:
            new_box = Box(box.name, Ty(), box.dom.r @ box.cod)
            return box.id(box.dom) @ new_box >> box.cups(box.dom, box.dom.r) @ box.id(box.cod) 
        else:
            new_box = Box(box.name, Ty(), box.cod @ box.dom.l)
            return new_box @ box.id(box.dom) >> box.id(box.cod) @ box.cups(box.dom.l, box.dom)
```

Do we want a `uncurry` method for `Diagram` to abstract away the uncurrying constructions?@dimkart: @kinianlo For the moment your uncurry method is fine I guess -- we can think later for a more clever distribution of the inputs along the boxes. Your uncurry rewrite rule is also fine for now, no need to add it to the Diagram (since if we do that, some other parts of the codebase need to be updated as well to use the new methods).

Solution patch:
diff --git a/lambeq/__init__.py b/lambeq/__init__.py
index 726fc01..372d838 100644
--- a/lambeq/__init__.py
+++ b/lambeq/__init__.py
@@ -67,6 +67,7 @@ __all__ = [
         'stairs_reader',
         'word_sequence_reader',
 
+        'CollapseDomainRewriteRule',
         'CoordinationRewriteRule',
         'CurryRewriteRule',
         'DiagramRewriter',
@@ -115,11 +116,11 @@ from lambeq.ansatz import (BaseAnsatz, CircuitAnsatz, IQPAnsatz, MPSAnsatz,
                            StronglyEntanglingAnsatz, TensorAnsatz)
 from lambeq.core.globals import VerbosityLevel
 from lambeq.core.types import AtomicType
-from lambeq.rewrite import (CoordinationRewriteRule, CurryRewriteRule,
-                            DiagramRewriter, RemoveCupsRewriter,
-                            RemoveSwapsRewriter, Rewriter, RewriteRule,
-                            SimpleRewriteRule, UnifyCodomainRewriter,
-                            UnknownWordsRewriteRule)
+from lambeq.rewrite import (CollapseDomainRewriteRule, CoordinationRewriteRule,
+                            CurryRewriteRule, DiagramRewriter,
+                            RemoveCupsRewriter, RemoveSwapsRewriter, Rewriter,
+                            RewriteRule, SimpleRewriteRule,
+                            UnifyCodomainRewriter, UnknownWordsRewriteRule)
 from lambeq.text2diagram import (
         CCGType, CCGRule, CCGRuleUseError, CCGTree,
         CCGParser,
diff --git a/lambeq/ansatz/tensor.py b/lambeq/ansatz/tensor.py
index 0d43f6b..09c8460 100644
--- a/lambeq/ansatz/tensor.py
+++ b/lambeq/ansatz/tensor.py
@@ -29,6 +29,7 @@ from lambeq.ansatz.base import AnsatzWithFramesRuntimeError, BaseAnsatz
 from lambeq.backend import grammar, Symbol, tensor
 from lambeq.backend.grammar import Cup, Spider, Ty, Word
 from lambeq.backend.tensor import Dim
+from lambeq.rewrite import CollapseDomainRewriteRule
 
 
 class TensorAnsatz(BaseAnsatz):
@@ -121,6 +122,25 @@ class SplitTensorAnsatz(TensorAnsatz):
     into smaller units."""
 
     split_functor: grammar.Functor
+    uncurry: CollapseDomainRewriteRule
+
+    def __init__(self, ob_map: Mapping[grammar.Ty, tensor.Dim],
+                 uncurry_left: bool) -> None:
+        """Instantiate a split tensor network ansatz.
+
+        Parameters
+        ----------
+        ob_map : dict
+            A mapping from :py:class:`lambeq.backend.grammar.Ty` to the
+            dimension space it uses in a tensor network.
+        uncurry_left: bool
+            If True, the uncurrying cups are placed on the left-hand
+            side. If False, they are placed on the right-hand side.
+
+        """
+
+        super().__init__(ob_map)
+        self.uncurry = CollapseDomainRewriteRule(left=uncurry_left)
 
     @abstractmethod
     def _split_ar(self, _: grammar.Functor, ar: Word) -> grammar.Diagrammable:
@@ -143,7 +163,8 @@ class MPSAnsatz(SplitTensorAnsatz):
     def __init__(self,
                  ob_map: Mapping[Ty, Dim],
                  bond_dim: int,
-                 max_order: int = 3) -> None:
+                 max_order: int = 3,
+                 uncurry_left: bool = True) -> None:
         """Instantiate a matrix product state ansatz.
 
         Parameters
@@ -156,6 +177,9 @@ class MPSAnsatz(SplitTensorAnsatz):
         max_order: int
             The maximum order of each tensor in the matrix product
             state, which must be at least 3.
+        uncurry_left: bool
+            If True, the uncurrying cups are placed on the left-hand
+            side. If False, they are placed on the right-hand side.
 
         """
         if max_order < 3:
@@ -165,21 +189,25 @@ class MPSAnsatz(SplitTensorAnsatz):
         ob_map = dict(ob_map)
         ob_map[self.BOND_TYPE] = Dim(bond_dim)
 
-        super().__init__(ob_map)
+        super().__init__(ob_map, uncurry_left)
 
         self.bond_dim = bond_dim
         self.max_order = max_order
         self.split_functor = grammar.Functor(
             grammar.grammar,
             ob=lambda _, ob: ob,
-            ar=self._split_ar  # type: ignore[arg-type]
+            ar=self._split_ar
         )
 
-    def _split_ar(self, _: grammar.Functor, ar: Word) -> grammar.Diagrammable:
-        bond = self.BOND_TYPE
-        if len(ar.cod) <= self.max_order:
-            return Word(f'{ar.name}_0', ar.cod)
+    def _split_ar(self, _: grammar.Functor,
+                  ar: grammar.Box) -> grammar.Diagrammable:
+        if len(ar.dom) + len(ar.cod) <= self.max_order:
+            return grammar.Box(f'{ar.name}_0', ar.dom, ar.cod, z=ar.z)
+
+        if self.uncurry.matches(ar):
+            return self.split_functor(self.uncurry.rewrite(ar))
 
+        bond = self.BOND_TYPE
         boxes = []
         cups = []
         step_size = self.max_order - 2
@@ -199,7 +227,8 @@ class SpiderAnsatz(SplitTensorAnsatz):
 
     def __init__(self,
                  ob_map: Mapping[Ty, Dim],
-                 max_order: int = 2) -> None:
+                 max_order: int = 2,
+                 uncurry_left: bool = True) -> None:
         """Instantiate a spider ansatz.
 
         Parameters
@@ -209,23 +238,30 @@ class SpiderAnsatz(SplitTensorAnsatz):
             dimension space it uses in a tensor network.
         max_order: int
             The maximum order of each tensor, which must be at least 2.
+        uncurry_left: bool
+            If True, the uncurrying cups are placed on the left-hand
+            side. If False, they are placed on the right-hand side.
 
         """
         if max_order < 2:
             raise ValueError('`max_order` must be at least 2')
 
-        super().__init__(ob_map)
+        super().__init__(ob_map, uncurry_left)
 
         self.max_order = max_order
         self.split_functor = grammar.Functor(
             grammar.grammar,
             ob=lambda _, ob: ob,
-            ar=self._split_ar  # type: ignore[arg-type]
+            ar=self._split_ar
         )
 
-    def _split_ar(self, _: grammar.Functor, ar: Word) -> grammar.Diagrammable:
-        if len(ar.cod) <= self.max_order:
-            return Word(f'{ar.name}_0', ar.cod)
+    def _split_ar(self, _: grammar.Functor,
+                  ar: grammar.Box) -> grammar.Diagrammable:
+        if len(ar.dom) + len(ar.cod) <= self.max_order:
+            return grammar.Box(f'{ar.name}_0', ar.dom, ar.cod, z=ar.z)
+
+        if self.uncurry.matches(ar):
+            return self.split_functor(self.uncurry.rewrite(ar))
 
         boxes = []
         spiders = [grammar.Id(ar.cod[:1])]
diff --git a/lambeq/rewrite/__init__.py b/lambeq/rewrite/__init__.py
index 5a9b04e..4654abe 100644
--- a/lambeq/rewrite/__init__.py
+++ b/lambeq/rewrite/__init__.py
@@ -12,7 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-__all__ = ['CoordinationRewriteRule',
+__all__ = ['CollapseDomainRewriteRule',
+           'CoordinationRewriteRule',
            'CurryRewriteRule',
            'Rewriter',
            'RewriteRule',
@@ -24,9 +25,10 @@ __all__ = ['CoordinationRewriteRule',
            'RemoveSwapsRewriter',
            'UnifyCodomainRewriter']
 
-from lambeq.rewrite.base import (CoordinationRewriteRule, CurryRewriteRule,
-                                 RewriteRule, Rewriter, SimpleRewriteRule,
-                                 UnknownWordsRewriteRule)
+from lambeq.rewrite.base import (CollapseDomainRewriteRule,
+                                 CoordinationRewriteRule,
+                                 CurryRewriteRule, RewriteRule, Rewriter,
+                                 SimpleRewriteRule, UnknownWordsRewriteRule)
 
 from lambeq.rewrite.rewrite_diagram import (DiagramRewriter,
                                             RemoveCupsRewriter,
diff --git a/lambeq/rewrite/base.py b/lambeq/rewrite/base.py
index a817063..2ef4783 100644
--- a/lambeq/rewrite/base.py
+++ b/lambeq/rewrite/base.py
@@ -321,6 +321,45 @@ class CurryRewriteRule(RewriteRule):
         return new_box
 
 
+class CollapseDomainRewriteRule(RewriteRule):
+    """
+    A rewrite rule that convert boxes into domain-less boxes
+    by uncurrying.
+    """
+    def __init__(self, left: bool = True) -> None:
+        """Instantiate a CollapseDomainRewriteRule.
+        This rules uncurries all the inputs of a box to create
+        a domain-less box. It can be used prior to applying an
+        ansatz that works only on domain-less boxes.
+
+        Parameters
+        ----------
+        left : bool, default: True
+            If True, the uncurrying cups are placed on the left-hand
+            side. If False, they are placed on the right-hand side.
+
+        """
+        self.left = left
+
+    def matches(self, box: Box) -> bool:
+        return bool(box.dom)
+
+    def rewrite(self, box: Box) -> Diagrammable:
+        Cup = box.category.Diagram.special_boxes['cup']
+        Id = box.category.Diagram.id
+
+        if self.left:
+            new_box = Box(box.name, Ty(), box.dom.r @ box.cod)
+            first_layer = Id(box.dom) @ new_box
+            second_layer = Cup(box.dom, box.dom.r) @ Id(box.cod)
+            return first_layer.then(second_layer)
+        else:
+            new_box = Box(box.name, Ty(), box.cod @ box.dom.l)
+            first_layer = new_box @ Id(box.dom)
+            second_layer = Id(box.cod) @ box.cups(box.dom.l, box.dom)
+            return first_layer.then(second_layer)
+
+
 class Rewriter:
     """Class that rewrites diagrams.
 


Test patch:
diff --git a/tests/ansatz/test_tensor.py b/tests/ansatz/test_tensor.py
index bc3f207..e739ea8 100644
--- a/tests/ansatz/test_tensor.py
+++ b/tests/ansatz/test_tensor.py
@@ -29,6 +29,13 @@ def diagram_with_frame():
     )
 
 
+@pytest.fixture
+def diagram_with_box():
+    cod = Ty(objects=list(map(Ty, 'abcd')))
+    return ((Word('big', cod.l) @ Word('words', cod @ Ty('a')))
+            >> Box('box', cod.l @ cod @ Ty('a'), cod @ cod))
+
+
 def test_tensor_ansatz_eval(diagram):
     ob_map = {Ty(t): Dim(2) for t in 'abcd'}
     ansatz = TensorAnsatz(ob_map)
@@ -56,6 +63,15 @@ def test_mps_ansatz_split(diagram):
             assert len(box.cod) <= i
 
 
+def test_mps_ansatz_split_with_box(diagram_with_box):
+    ob_map = {Ty(t): Dim(4) for t in 'abcd'}
+    for i in range(3, 5):
+        ansatz = MPSAnsatz(ob_map, max_order=i, bond_dim=1)
+        split_diagram = ansatz(diagram_with_box)
+        for box in split_diagram.boxes:
+            assert len(box.dom) + len(box.cod) <= i
+
+
 def test_mps_ansatz_eval(diagram):
     ob_map = {Ty(t): Dim(4) for t in 'abcd'}
     ansatz = MPSAnsatz(ob_map, bond_dim=1)
@@ -80,6 +96,18 @@ def test_spider_splitter(diagram):
             assert len(box.cod) <= i
 
 
+def test_spider_splitter_with_box(diagram_with_box):
+    from lambeq.backend.grammar import Spider
+    ob_map = {Ty(t): Dim(4) for t in 'abcd'}
+
+    for i in range(2, 5):
+        ansatz = SpiderAnsatz(ob_map, max_order=i)
+        tensor = ansatz(diagram_with_box)
+        for box in tensor.boxes:
+            if not isinstance(box, Spider):
+                assert len(box.dom) + len(box.cod) <= i
+
+
 def test_spider_ansatz_eval(diagram):
     ob_map = {Ty(t): Dim(4) for t in 'abcd'}
     ansatz = SpiderAnsatz(ob_map)
diff --git a/tests/test_rewrite.py b/tests/test_rewrite.py
index 905d43a..6e6d781 100644
--- a/tests/test_rewrite.py
+++ b/tests/test_rewrite.py
@@ -6,7 +6,7 @@ from lambeq.backend.grammar import (Box, Cap, Cup, Diagram, Id, Spider,
 from lambeq import (AtomicType, Rewriter, CoordinationRewriteRule,
                     CurryRewriteRule, RemoveCupsRewriter,
                     RemoveSwapsRewriter, SimpleRewriteRule, stairs_reader,
-                    UnifyCodomainRewriter, UnknownWordsRewriteRule)
+                    CollapseDomainRewriteRule, UnifyCodomainRewriter, UnknownWordsRewriteRule)
 
 N = AtomicType.NOUN
 S = AtomicType.SENTENCE
@@ -466,3 +466,18 @@ def test_unknown_words_stairs():
     ])
 
     assert rewrite_unknown_words(diagram) == diagram
+
+
+def test_collapse_domain():
+    uncurry = CollapseDomainRewriteRule()
+    dom = Ty(objects=list(map(Ty, 'abcd')))
+    cod = Ty(objects=list(map(Ty, 'wxyz')))
+    box = Box('box', dom, cod)
+    assert uncurry.matches(box)
+    diag = uncurry.rewrite(box)
+
+    assert diag.dom == box.dom
+    assert diag.cod == box.cod
+
+    num_cups = sum(1 for box in diag.boxes if isinstance(box, Cup))
+    assert num_cups == len(dom)