Adding backend interface and numpy backend to Mr Mustard

.github/workflows/tests_numpy.yml

@@ -0,0 +1,47 @@
+name: Numpy tests
+on:
+  push:
+    branches:
+      - develop
+  pull_request:
+    paths:
+      - '.github/workflows/tests.yml'
+      - 'mrmustard/**'
+      - 'tests/**'
+      - 'pyproject.toml'
+      - 'poetry.lock'
+      - 'pytest.ini'
+
+jobs:
+  pytest:
+    runs-on: ubuntu-latest
+    concurrency:
+      group: ${{ github.workflow }}-${{ github.ref }}
+      cancel-in-progress: true
+    env:
+      HYPOTHESIS_PROFILE: ci
+
+    steps:
+      - name: Checkout repo
+        uses: actions/checkout@v3
+
+      - name: Setup python
+        uses: actions/setup-python@v4
+        with:
+          python-version: '3.9'
+
+      - name: Setup Julia
+        uses: julia-actions/setup-julia@v1
+        with:
+          version: 1.9.3
+
+      - name: Install dependencies
+        run: |
+          python -m pip install --no-cache-dir --upgrade pip
+          pip install --no-cache-dir poetry==1.4.0
+          poetry config virtualenvs.create false
+          poetry install --extras "ray" --with dev
+          julia --project="julia_pkg" -e "using Pkg; Pkg.instantiate()"
+
+      - name: Run tests
+        run: python -m pytest tests --backend=numpy -p no:warnings --tb=native

.github/workflows/tests.yml → .github/workflows/tests_tensorflow.yml

@@ -1,4 +1,4 @@
-name: Tests
+name: Tensorflow Tests
 on:
   push:
     branches:
@@ -44,7 +44,7 @@ jobs:
           julia --project="julia_pkg" -e "using Pkg; Pkg.instantiate()"
 
       - name: Run tests
-        run: python -m pytest tests --cov=mrmustard --cov-report=term-missing --cov-report=xml -p no:warnings --tb=native
+        run: python -m pytest tests --backend=tensorflow --cov=mrmustard --cov-report=term-missing --cov-report=xml -p no:warnings --tb=native
 
       - name: Upload coverage to Codecov
         uses: codecov/codecov-action@v3

README.md

@@ -235,8 +235,7 @@ The physics module contains a growing number of functions that we can apply to s
 The math module is the backbone of Mr Mustard, which consists in the [`Math`](https://github.com/XanaduAI/MrMustard/blob/main/mrmustard/math/math_interface.py) interface. Mr Mustard comes with a plug-and-play backends through a math interface. You can use it as a drop-in replacement for tensorflow or pytorch and your code will be plug-and-play too!
 ```python
 from mrmustard import settings
-from mrmustard.math import Math
-math = Math()
+import mrmustard.math as math
 
 math.cos(0.1)  # tensorflow
 

doc/code/math.rst

@@ -3,15 +3,15 @@ mrmustard.math
 
 .. currentmodule:: mrmustard.math
 
-Backends
----------
+The `math` module
+-----------------
 
 .. toctree::
     :maxdepth: 1
 
     math/parameter_set
     math/parameter
-    math/tensorflow
+    math/backend_manager
 
 .. automodapi:: mrmustard.math
     :no-heading:

doc/code/math/backend_manager.rst

@@ -0,0 +1,8 @@
+Backend Manager
+===============
+
+.. currentmodule:: mrmustard.math.backend_manager
+
+.. automodapi:: mrmustard.math.backend_manager
+    :no-heading:
+    :include-all-objects:

doc/code/math/backend_numpy.rst

@@ -0,0 +1,8 @@
+Backend Manager
+===============
+
+.. currentmodule:: mrmustard.math.backend_numpy
+
+.. automodapi:: mrmustard.math.backend_numpy
+    :no-heading:
+    :include-all-objects:

doc/introduction/basic_reference.md

@@ -182,8 +182,7 @@ The math module is the backbone of Mr Mustard, which consists in the [Math](http
 Mr Mustard comes with a plug-and-play backends through a math interface. You can use it as a drop-in replacement for tensorflow or pytorch and your code will be plug-and-play too!
 ```python
 from mrmustard import settings
-from mrmustard.math import Math
-math = Math()
+import mrmustard.math as math
 
 math.cos(0.1)  # tensorflow
 

mrmustard/lab/abstract/measurement.py

@@ -20,14 +20,11 @@
 from typing import Iterable, Sequence, Union
 
 from mrmustard import settings
-from mrmustard.math import Math
 from mrmustard.math.parameter_set import ParameterSet
 from mrmustard.math.parameters import Constant, Variable
 from mrmustard.utils.typing import Tensor
-
 from .state import State
-
-math = Math()
+import mrmustard.math as math
 
 
 class Measurement(ABC):

mrmustard/lab/abstract/state.py

@@ -32,7 +32,6 @@
 from matplotlib import cm
 
 from mrmustard import settings
-from mrmustard.math import Math
 from mrmustard.math.parameters import Constant, Variable
 from mrmustard.physics import bargmann, fock, gaussian
 from mrmustard.utils.typing import (
@@ -44,13 +43,12 @@
     RealVector,
 )
 from mrmustard.physics.wigner import wigner_discretized
+import mrmustard.math as math
 
 
 if TYPE_CHECKING:
     from .transformation import Transformation
 
-math = Math()
-
 
 # pylint: disable=too-many-instance-attributes
 class State:  # pylint: disable=too-many-public-methods

mrmustard/lab/abstract/transformation.py

@@ -24,7 +24,6 @@
 import numpy as np
 
 from mrmustard import settings
-from mrmustard.math import Math
 from mrmustard.math.parameter_set import ParameterSet
 from mrmustard.math.parameters import Constant, Variable
 from mrmustard.math.tensor_networks import Tensor
@@ -33,7 +32,7 @@
 
 from .state import State
 
-math = Math()
+import mrmustard.math as math
 
 
 class Transformation(Tensor):
@@ -122,9 +121,6 @@ def _transform_gaussian(self, state: State, dual: bool) -> State:
         )  # NOTE: assumes modes don't change
         return new_state
 
-    def _transform_fock(self, state: State, dual: bool) -> State:
-        raise NotImplementedError
-
     @property
     def num_modes(self) -> int:
         r"""The number of modes on which the transformation acts."""

mrmustard/lab/circuit.py

@@ -26,7 +26,7 @@
 from mrmustard.lab.abstract import State, Transformation
 from mrmustard.utils.typing import RealMatrix, RealVector
 from mrmustard.lab.circuit_drawer import circuit_text
-from mrmustard.math.xptensor import XPMatrix, XPVector
+from mrmustard.math.tensor_wrappers import XPMatrix, XPVector
 
 import numpy as np
 

mrmustard/lab/detectors.py

@@ -19,7 +19,6 @@
 from typing import Iterable, List, Optional, Tuple, Union
 
 from mrmustard import settings
-from mrmustard.math import Math
 from mrmustard.physics import fock, gaussian
 from mrmustard.utils.typing import RealMatrix, RealVector
 
@@ -28,7 +27,7 @@
 from .states import Coherent, DisplacedSqueezed
 from .utils import make_parameter
 
-math = Math()
+import mrmustard.math as math
 
 __all__ = ["PNRDetector", "ThresholdDetector", "Generaldyne", "Homodyne", "Heterodyne"]
 

mrmustard/lab/gates.py

@@ -22,14 +22,12 @@
 import numpy as np
 
 from mrmustard import settings
-from mrmustard.math import Math
-from mrmustard.math.parameters import update_orthogonal, update_symplectic, update_unitary
-from mrmustard.physics import fock, gaussian
+from mrmustard.physics import gaussian, fock
 from mrmustard.utils.typing import ComplexMatrix, RealMatrix
+from mrmustard.math.parameters import update_orthogonal, update_symplectic, update_unitary
 from .abstract import Channel, Unitary, State
 from .utils import make_parameter
-
-math = Math()
+import mrmustard.math as math
 
 __all__ = [
     "Dgate",
@@ -114,6 +112,8 @@ def U(self, cutoffs: Optional[Sequence[int]] = None, shape: Optional[Sequence[in
            Raises:
                ValueError: if the length of the cutoffs array is different from N and 2N
         """
+        from mrmustard.physics.fock_custom_grads import displacement
+
         N = self.num_modes
         if cutoffs is None:
             pass
@@ -137,17 +137,17 @@ def U(self, cutoffs: Optional[Sequence[int]] = None, shape: Optional[Sequence[in
             Ud = None
             for idx, out_in in enumerate(zip(shape[:N], shape[N:])):
                 if Ud is None:
-                    Ud = fock.displacement(x[idx], y[idx], shape=out_in)
+                    Ud = displacement(x[idx], y[idx], shape=out_in)
                 else:
-                    U_next = fock.displacement(x[idx], y[idx], shape=out_in)
+                    U_next = displacement(x[idx], y[idx], shape=out_in)
                     Ud = math.outer(Ud, U_next)
 
             return math.transpose(
                 Ud,
                 list(range(0, 2 * N, 2)) + list(range(1, 2 * N, 2)),
             )
         else:
-            return fock.displacement(x[0], y[0], shape=shape)
+            return displacement(x[0], y[0], shape=shape)
 
 
 class Sgate(Unitary):
@@ -206,6 +206,8 @@ def U(self, cutoffs: Optional[Sequence[int]] = None, shape: Optional[Sequence[in
         Returns:
             array[complex]: the unitary matrix
         """
+        from mrmustard.physics.fock_custom_grads import squeezer
+
         N = self.num_modes
         if cutoffs is None:
             pass
@@ -230,16 +232,16 @@ def U(self, cutoffs: Optional[Sequence[int]] = None, shape: Optional[Sequence[in
             Us = None
             for idx, single_shape in enumerate(zip(shape[:N], shape[N:])):
                 if Us is None:
-                    Us = fock.squeezer(r[idx], phi[idx], shape=single_shape)
+                    Us = squeezer(r[idx], phi[idx], shape=single_shape)
                 else:
-                    U_next = fock.squeezer(r[idx], phi[idx], shape=single_shape)
+                    U_next = squeezer(r[idx], phi[idx], shape=single_shape)
                     Us = math.outer(Us, U_next)
             return math.transpose(
                 Us,
                 list(range(0, 2 * N, 2)) + list(range(1, 2 * N, 2)),
             )
         else:
-            return fock.squeezer(r[0], phi[0], shape=shape)
+            return squeezer(r[0], phi[0], shape=shape)
 
     @property
     def X_matrix(self):
@@ -511,6 +513,8 @@ def U(
         Returns:
             array[complex]: the unitary tensor of the beamsplitter
         """
+        from mrmustard.physics.fock_custom_grads import beamsplitter
+
         if cutoffs is None:
             pass
         elif len(cutoffs) == 4:
@@ -522,7 +526,7 @@ def U(
 
         shape = shape or cutoffs
 
-        return fock.beamsplitter(
+        return beamsplitter(
             self.theta.value,
             self.phi.value,
             shape=shape,

mrmustard/lab/states.py

@@ -19,7 +19,6 @@
 from typing import List, Optional, Sequence, Tuple, Union
 
 from mrmustard import settings
-from mrmustard.math import Math
 from mrmustard.math.parameters import update_symplectic
 from mrmustard.math.parameter_set import ParameterSet
 from mrmustard.physics import fock, gaussian
@@ -28,7 +27,7 @@
 from .abstract import State
 from .utils import make_parameter
 
-math = Math()
+import mrmustard.math as math
 
 
 __all__ = [

mrmustard/lab/utils.py

@@ -20,10 +20,8 @@
 from typing import Callable, Optional, Tuple
 from mrmustard.math.parameters import update_euclidean
 
-from mrmustard.math import Math
 from mrmustard.math.parameters import Constant, Variable
-
-math = Math()
+import mrmustard.math as math
 
 
 def make_parameter(