Release version 0.2.8

Co-authored-by: Thomas Hoffmann <thomas.hoffmann@cambridgequantum.com>
Co-authored-by: Nikhil Khatri <nikhil.khatri@quantinuum.com>

.github/workflows/build_test.yml

@@ -21,9 +21,9 @@ jobs:
     outputs:
       error-check: ${{ steps.error-check.conclusion }}
     steps:
-    - uses: actions/checkout@v2
+    - uses: actions/checkout@v3
     - name: Setup Python ${{ matrix.python-version }}
-      uses: actions/setup-python@v2
+      uses: actions/setup-python@v4
       with:
         python-version: ${{ matrix.python-version }}
     - name: Install linter
@@ -49,25 +49,14 @@ jobs:
       matrix:
         python-version: [ 3.8, 3.9, "3.10" ]
     steps:
-    - uses: actions/checkout@v2
+    - uses: actions/checkout@v3
     - name: Setup Python ${{ matrix.python-version }}
-      uses: actions/setup-python@v2
+      uses: actions/setup-python@v4
       with:
         python-version: ${{ matrix.python-version }}
-    - name: Locate pip cache
-      id: loc-pip-cache
-      run: echo "::set-output name=dir::$(pip cache dir)"
-    - name: Restore pip dependencies from cache
-      uses: actions/cache@v2
-      with:
-        path: ${{ steps.loc-pip-cache.outputs.dir }}
-        key: build_and_test-${{ runner.os }}-pip-${{ matrix.python-version }}-${{ hashFiles('setup.cfg') }}
-        restore-keys: |
-          build_and_test-${{ runner.os }}-pip-${{ matrix.python-version }}-
-          build_and_test-${{ runner.os }}-pip-
-    - name: Install DisCoPy from GitHub
+    - name: Install DisCoPy 0.5 from GitHub
       if: github.ref_name != 'release' && github.ref_name != 'beta'
-      run: pip install git+https://github.com/oxford-quantum-group/discopy
+      run: pip install git+https://github.com/discopy/discopy@0.5
     - name: Install base package
       run: pip install .
     - name: Check package import works
@@ -130,9 +119,9 @@ jobs:
       matrix:
         python-version: [ 3.8, 3.9, "3.10" ]
     steps:
-    - uses: actions/checkout@v2
+    - uses: actions/checkout@v3
     - name: Setup Python ${{ matrix.python-version }}
-      uses: actions/setup-python@v2
+      uses: actions/setup-python@v4
       with:
         python-version: ${{ matrix.python-version }}
     - name: Install dependencies with type hints

.github/workflows/docs.yml

@@ -16,11 +16,11 @@ jobs:
     name: Build and deploy documentation
     runs-on: ubuntu-latest
     steps:
-    - uses: actions/checkout@v2
+    - uses: actions/checkout@v3
       with:
         fetch-depth: 0  # fetches tags, required for version info
     - name: Set up Python
-      uses: actions/setup-python@v2
+      uses: actions/setup-python@v4
       with:
         python-version: 3.8
     - name: Build lambeq

docs/conf.py

@@ -46,7 +46,7 @@
 ]
 
 intersphinx_mapping = {
-    'discopy': ("https://discopy.readthedocs.io/en/main/", None),
+    'discopy': ("https://discopy.readthedocs.io/en/0.5/", None),
     'pennylane': ("https://pennylane.readthedocs.io/en/stable/", None),
 }
 

docs/release_notes.rst

@@ -3,6 +3,20 @@
 Release notes
 =============
 
+.. _rel-0.2.8:
+
+`0.2.8 <https://github.com/CQCL/lambeq/releases/tag/0.2.8>`_
+------------------------------------------------------------
+Changed:
+
+- Improved the performance of :py:class:`.NumpyModel` when using Jax JIT-compilation.
+- Dependencies: pinned the required version of DisCoPy to 0.5.X.
+
+Fixed:
+
+- Fixed incorrectly scaled validation loss in progress bar during model training.
+- Fixed symbol type mismatch in the quantum models when a circuit was previously converted to tket.
+
 .. _rel-0.2.7:
 
 `0.2.7 <https://github.com/CQCL/lambeq/releases/tag/0.2.7>`_

lambeq/__init__.py

@@ -17,11 +17,10 @@
         '__version_info__',
 
         'ansatz',
-        'text2diagram',
         'core',
         'pregroups',
-        'reader',
         'rewrite',
+        'text2diagram',
         'tokeniser',
         'training',
 

lambeq/ansatz/circuit.py

@@ -34,9 +34,9 @@
 from discopy.quantum.gates import Bra, H, Ket, Rx, Ry, Rz
 from discopy.rigid import Box, Diagram, Ty
 import numpy as np
-from sympy import symbols
+from sympy import Symbol, symbols
 
-from lambeq.ansatz import BaseAnsatz, Symbol
+from lambeq.ansatz import BaseAnsatz
 
 computational_basis = Id(qubit)
 

lambeq/text2diagram/ccg_parser.py

@@ -22,7 +22,7 @@
 from typing import Optional
 
 from discopy import Diagram
-from tqdm.autonotebook import tqdm
+from tqdm.auto import tqdm
 
 from lambeq.core.globals import VerbosityLevel
 from lambeq.core.utils import (SentenceBatchType, SentenceType,

lambeq/training/model.py

@@ -26,7 +26,7 @@
 from typing import Any, Union
 
 from discopy.tensor import Diagram
-from sympy import default_sort_key
+from sympy import default_sort_key, Symbol as SymPySymbol
 
 from lambeq.ansatz.base import Symbol
 from lambeq.training.checkpoint import Checkpoint
@@ -50,7 +50,7 @@ class Model(ABC):
 
     def __init__(self) -> None:
         """Initialise an instance of :py:class:`Model` base class."""
-        self.symbols: list[Symbol] = []
+        self.symbols: list[Union[Symbol, SymPySymbol]] = []
         self.weights: Collection = []
 
     def __call__(self, *args: Any, **kwds: Any) -> Any:

lambeq/training/numpy_model.py

@@ -28,14 +28,19 @@
 
 from collections.abc import Callable, Iterable
 import pickle
-from typing import Any
+from typing import Any, TYPE_CHECKING, Union
 
 from discopy import Tensor
 from discopy.tensor import Diagram
 import numpy
 from numpy.typing import ArrayLike
 from sympy import lambdify
 
+
+if TYPE_CHECKING:
+    from jax import numpy as jnp
+
+
 from lambeq.training.quantum_model import QuantumModel
 
 
@@ -74,7 +79,7 @@ def _get_lambda(self, diagram: Diagram) -> Callable[[Any], Any]:
         if diagram in self.lambdas:
             return self.lambdas[diagram]
 
-        def diagram_output(*x: ArrayLike) -> ArrayLike:
+        def diagram_output(x: Iterable[ArrayLike]) -> ArrayLike:
             with Tensor.backend('jax'), tn.DefaultBackend('jax'):
                 sub_circuit = self._fast_subs([diagram], x)[0]
                 result = tn.contractors.auto(*sub_circuit.to_tn()).tensor
@@ -112,7 +117,9 @@ def _fast_subs(self,
                         b._phase = b._data
         return diagrams
 
-    def get_diagram_output(self, diagrams: list[Diagram]) -> numpy.ndarray:
+    def get_diagram_output(self,
+                           diagrams: list[Diagram]) -> Union[jnp.ndarray,
+                                                             numpy.ndarray]:
         """Return the exact prediction for each diagram.
 
         Parameters
@@ -142,9 +149,13 @@ def get_diagram_output(self, diagrams: list[Diagram]) -> numpy.ndarray:
                              'from pre-trained checkpoint.')
 
         if self.use_jit:
+            from jax import numpy as jnp
+
             lambdified_diagrams = [self._get_lambda(d) for d in diagrams]
-            return numpy.array([diag_f(*self.weights)
-                                for diag_f in lambdified_diagrams])
+            res: jnp.ndarray = jnp.array([diag_f(self.weights)
+                                          for diag_f in lambdified_diagrams])
+
+            return res
 
         diagrams = self._fast_subs(diagrams, self.weights)
         with Tensor.backend('numpy'):

lambeq/training/pytorch_model.py

@@ -35,6 +35,7 @@ class PytorchModel(Model, torch.nn.Module):
     """A lambeq model for the classical pipeline using PyTorch."""
 
     weights: torch.nn.ParameterList  # type: ignore[assignment]
+    symbols: list[Symbol]  # type: ignore[assignment]
 
     def __init__(self) -> None:
         """Initialise a PytorchModel."""

lambeq/training/quantum_model.py

@@ -21,11 +21,16 @@
 from __future__ import annotations
 
 from abc import abstractmethod
-from typing import Any
+from typing import Any, TYPE_CHECKING, Union
 
 from discopy.tensor import Diagram, Tensor
 import numpy as np
 
+
+if TYPE_CHECKING:
+    from jax import numpy as jnp
+
+
 from lambeq.training.checkpoint import Checkpoint
 from lambeq.training.model import Model
 
@@ -90,7 +95,6 @@ def initialise_weights(self) -> None:
         if not self.symbols:
             raise ValueError('Symbols not initialised. Instantiate through '
                              '`from_diagrams()`.')
-        assert all(w.size == 1 for w in self.symbols)
         self.weights = np.random.rand(len(self.symbols))
 
     def _load_checkpoint(self, checkpoint: Checkpoint) -> None:
@@ -124,7 +128,8 @@ def _make_checkpoint(self) -> Checkpoint:
         return checkpoint
 
     @abstractmethod
-    def get_diagram_output(self, diagrams: list[Diagram]) -> np.ndarray:
+    def get_diagram_output(self, diagrams: list[Diagram]) -> Union[jnp.ndarray,
+                                                                   np.ndarray]:
         """Return the diagram prediction.
 
         Parameters

lambeq/training/trainer.py

@@ -411,6 +411,7 @@ def writer_helper(*args: Any) -> None:
             if val_dataset is not None:
                 if epoch % evaluation_step == 0:
                     val_loss = 0.0
+                    seen_so_far = 0
                     batches_per_validation = ceil(len(val_dataset)
                                                   / val_dataset.batch_size)
                     with Tensor.backend(self.backend):
@@ -425,6 +426,7 @@ def writer_helper(*args: Any) -> None:
                             x_val, y_label_val = v_batch
                             y_hat_val, cur_loss = self.validation_step(v_batch)
                             val_loss += cur_loss * len(x_val)
+                            seen_so_far += len(x_val)
                             if self.evaluate_functions is not None:
                                 for metr, func in (
                                         self.evaluate_functions.items()):
@@ -434,7 +436,7 @@ def writer_helper(*args: Any) -> None:
                             status_bar.set_description(
                                     self._generate_stat_report(
                                         train_loss=train_loss,
-                                        val_loss=val_loss))
+                                        val_loss=val_loss/seen_so_far))
                         val_loss /= len(val_dataset)
                         self.val_costs.append(val_loss)
                         status_bar.set_description(

setup.cfg

@@ -52,7 +52,7 @@ packages =
     lambeq.tokeniser
     lambeq.training
 install_requires =
-    discopy >= 0.4.3
+    discopy ~= 0.5.0
     pytket >= 0.19.2
     pyyaml
     spacy >= 3.0

tests/test_circuit.py

@@ -4,10 +4,11 @@
 from discopy.quantum import (Bra, CRz, CRx, CX, X, H, Ket,
                              qubit, Rx, Ry, Rz, sqrt, Controlled)
 from discopy.quantum.circuit import Circuit, Id
+from discopy.quantum.tk import from_tk
+from sympy import Symbol as sym
 
 from lambeq import (AtomicType, IQPAnsatz, Sim14Ansatz, Sim15Ansatz,
                     StronglyEntanglingAnsatz)
-from lambeq import Symbol as sym
 
 N = AtomicType.NOUN
 S = AtomicType.SENTENCE
@@ -255,3 +256,11 @@ def test_strongly_entangling_ansatz_ranges_error2():
     with pytest.raises(ValueError):
         ansatz = StronglyEntanglingAnsatz({q: 2}, 3, ranges=[1, 1, 2])
         ansatz(box)
+
+def test_discopy_tket_conversion():
+    word1, word2 = Word('Alice', N), Word('Bob', N.r)
+    sentence = word1 @ word2 >> Cup(N, N.r)
+    ansatz = IQPAnsatz({N: 1}, n_layers=1)
+    circuit = ansatz(sentence)
+    circuit_converted = from_tk(circuit.to_tk())
+    assert circuit.free_symbols == circuit_converted.free_symbols