Repository: XanaduAI/strawberryfields

Issue:
Bug in mean photon number
#### Issue description

Here is a minimal example that reproduces the bug:

```
import strawberryfields as sf
from strawberryfields import ops

modes = 4
fock_cutoff = 4
prog = sf.Program(modes)

with prog.context as q:
    ops.Sgate(.54) | q[0]
    ops.Sgate(.593) | q[1]
    ops.Sgate(.23) | q[2]
    ops.Sgate(.89) | q[3]

eng = sf.Engine("tf", backend_options={"cutoff_dim": fock_cutoff})
result = eng.run(prog)    
state = result.state
for i in range(0,modes):
    print("mean photon of mode " + str(i) + " = " + str((state.mean_photon(i)[0].numpy())))
```

This produces the output:

```
mean photon of mode 0 = 0.17980294
mean photon of mode 1 = 0.20575231
mean photon of mode 2 = 0.20575231
mean photon of mode 3 = 0.20575231
```

The mean photon of mode 2 and 3 should be different from mode 1 because the squeezing parameters are different.

* *Expected behavior:* (What you expect to happen)
```
mean photon of mode 0 = 0.17980294
mean photon of mode 1 = 0.20575231
mean photon of mode 2 = some different number
mean photon of mode 3 = some different number
```
* *Actual behavior:* (What actually happens)
```
mean photon of mode 0 = 0.17980294
mean photon of mode 1 = 0.20575231
mean photon of mode 2 = 0.20575231
mean photon of mode 3 = 0.20575231
```
* *Reproduces how often:* (What percentage of the time does it reproduce?)
```
100%
```
* *System information:* (post the output of `import strawberryfields as sf; sf.about()`)
```
Strawberry Fields: a Python library for continuous-variable quantum circuits.
Copyright 2018-2020 Xanadu Quantum Technologies Inc.

Python version:            3.8.2
Platform info:             Linux-5.6.8-arch1-1-x86_64-with-glibc2.2.5
Installation path:         /home/tvincent/.local/lib/python3.8/site-packages/strawberryfields
Strawberry Fields version: 0.15.1
Numpy version:             1.18.4
Scipy version:             1.4.1
SymPy version:             1.6
NetworkX version:          2.4
The Walrus version:        0.13.0
Blackbird version:         0.2.4
TensorFlow version:        2.3.1
```


Discussion:
Thanks for discovering and reporting this @trevor-vincent!

I've determined that the source of this bug is the function `tfbackend/ops.py:reduced_density_matrix` (used internally when computing `mean_photon`). The final `for` loop in this function fails to account for the fact that removing subsystems in `reduced_state` will cause the mode indices `m` to be "off by one" for each subsystem that has been previously removed during the loop. 

This is a bug at a core part of the codebase, so is high priority to fix. Thankfully, the fix looks relatively simple, something along the lines of:
```
for m in range(num_modes):
    if m != mode:
        reduced_state = partial_trace(reduced_state, m - removed_cnt, False, batched)
        removed_cnt += 1
```

Resolving patch:
diff --git a/strawberryfields/backends/tfbackend/ops.py b/strawberryfields/backends/tfbackend/ops.py
--- a/strawberryfields/backends/tfbackend/ops.py
+++ b/strawberryfields/backends/tfbackend/ops.py
@@ -62,7 +62,6 @@
     generate_bs_factors,
     load_bs_factors,
     save_bs_factors,
-    squeeze_parity,
 )
 
 def_type = tf.complex64  # NOTE: what if a user wants higher accuracy?
@@ -86,24 +85,27 @@ def mixed(pure_state, batched=False):
     # todo: In the fock backend mixing is done by ops.mix(), maybe the functions should be named identically?
     if not batched:
         pure_state = tf.expand_dims(pure_state, 0)  # add in fake batch dimension
+
     batch_offset = 1
     num_modes = len(pure_state.shape) - batch_offset
     max_num = (max_num_indices - batch_offset) // 2
+
     if num_modes > max_num:
         raise ValueError(
             "Converting state from pure to mixed currently only supported for {} modes.".format(
                 max_num
             )
         )
-    else:
-        # eqn: 'abc...xyz,ABC...XYZ->aAbBcC...xXyYzZ' (lowercase belonging to 'bra' side, uppercase belonging to 'ket' side)
-        batch_index = indices[:batch_offset]
-        bra_indices = indices[batch_offset : batch_offset + num_modes]
-        ket_indices = indices[batch_offset + num_modes : batch_offset + 2 * num_modes]
-        eqn_lhs = batch_index + bra_indices + "," + batch_index + ket_indices
-        eqn_rhs = "".join(bdx + kdx for bdx, kdx in zip(bra_indices, ket_indices))
-        eqn = eqn_lhs + "->" + batch_index + eqn_rhs
-        mixed_state = tf.einsum(eqn, pure_state, tf.math.conj(pure_state))
+
+    # eqn: 'abc...xyz,ABC...XYZ->aAbBcC...xXyYzZ' (lowercase belonging to 'bra' side, uppercase belonging to 'ket' side)
+    batch_index = indices[:batch_offset]
+    bra_indices = indices[batch_offset : batch_offset + num_modes]
+    ket_indices = indices[batch_offset + num_modes : batch_offset + 2 * num_modes]
+    eqn_lhs = batch_index + bra_indices + "," + batch_index + ket_indices
+    eqn_rhs = "".join(bdx + kdx for bdx, kdx in zip(bra_indices, ket_indices))
+    eqn = eqn_lhs + "->" + batch_index + eqn_rhs
+    mixed_state = tf.einsum(eqn, pure_state, tf.math.conj(pure_state))
+
     if not batched:
         mixed_state = tf.squeeze(mixed_state, 0)  # drop fake batch dimension
     return mixed_state
@@ -415,7 +417,7 @@ def fock_state(n, cutoff, pure=True, batched=False):
         raise ValueError("'n' is expected to be either an int or a numpy array")
     if batched:
         batch_size = n.shape[0]
-        idxs = [(b, f) for (b, f) in zip(range(batch_size), n)]
+        idxs = list(zip(range(batch_size), n))
         values = [1.0] * batch_size
         shape = [batch_size, cutoff]
     else:
@@ -537,45 +539,45 @@ def single_mode_gate(matrix, mode, in_modes, pure=True, batched=False):
         )
     if mode < 0 or mode >= num_modes:
         raise ValueError("'mode' argument is not compatible with number of in_modes")
-    else:
-        other_modes_indices = indices[
-            batch_offset + 2 * mode_size : batch_offset + (1 + num_modes) * mode_size
-        ]
-        if pure:
-            eqn_lhs = "{},{}{}{}{}".format(
-                batch_index + left_gate_str,
+
+    other_modes_indices = indices[
+        batch_offset + 2 * mode_size : batch_offset + (1 + num_modes) * mode_size
+    ]
+    if pure:
+        eqn_lhs = "{},{}{}{}{}".format(
+            batch_index + left_gate_str,
+            batch_index,
+            other_modes_indices[: mode * mode_size],
+            left_gate_str[1],
+            other_modes_indices[mode * mode_size :],
+        )
+        eqn_rhs = "".join(
+            [
                 batch_index,
                 other_modes_indices[: mode * mode_size],
-                left_gate_str[1],
+                left_gate_str[0],
                 other_modes_indices[mode * mode_size :],
-            )
-            eqn_rhs = "".join(
-                [
-                    batch_index,
-                    other_modes_indices[: mode * mode_size],
-                    left_gate_str[0],
-                    other_modes_indices[mode * mode_size :],
-                ]
-            )
-        else:
-            eqn_lhs = "{},{}{}{}{}{},{}".format(
-                batch_index + left_gate_str,
+            ]
+        )
+    else:
+        eqn_lhs = "{},{}{}{}{}{},{}".format(
+            batch_index + left_gate_str,
+            batch_index,
+            other_modes_indices[: mode * mode_size],
+            left_gate_str[1],
+            right_gate_str[0],
+            other_modes_indices[mode * mode_size :],
+            batch_index + right_gate_str,
+        )
+        eqn_rhs = "".join(
+            [
                 batch_index,
                 other_modes_indices[: mode * mode_size],
-                left_gate_str[1],
-                right_gate_str[0],
+                left_gate_str[0],
+                right_gate_str[1],
                 other_modes_indices[mode * mode_size :],
-                batch_index + right_gate_str,
-            )
-            eqn_rhs = "".join(
-                [
-                    batch_index,
-                    other_modes_indices[: mode * mode_size],
-                    left_gate_str[0],
-                    right_gate_str[1],
-                    other_modes_indices[mode * mode_size :],
-                ]
-            )
+            ]
+        )
 
     eqn = eqn_lhs + "->" + eqn_rhs
     einsum_inputs = [matrix, in_modes]
@@ -614,71 +616,68 @@ def two_mode_gate(matrix, mode1, mode2, in_modes, pure=True, batched=False):
         raise NotImplementedError(
             "The max number of supported modes for this operation is currently {}".format(max_len)
         )
+
+    min_mode = min(mode1, mode2)
+    max_mode = max(mode1, mode2)
+    if min_mode < 0 or max_mode >= num_modes or mode1 == mode2:
+        raise ValueError("One or more mode numbers are incompatible")
+
+    other_modes_indices = indices[
+        batch_offset + 4 * mode_size : batch_offset + 4 * mode_size + mode_size * (num_modes - 2)
+    ]
+    # build equation
+    if mode1 == min_mode:
+        lhs_min_mode_indices = left_gate_str[1]
+        lhs_max_mode_indices = left_gate_str[3]
+        rhs_min_mode_indices = left_gate_str[0]
+        rhs_max_mode_indices = left_gate_str[2]
     else:
-        min_mode = min(mode1, mode2)
-        max_mode = max(mode1, mode2)
-        if min_mode < 0 or max_mode >= num_modes or mode1 == mode2:
-            raise ValueError("One or more mode numbers are incompatible")
+        lhs_min_mode_indices = left_gate_str[3]
+        lhs_max_mode_indices = left_gate_str[1]
+        rhs_min_mode_indices = left_gate_str[2]
+        rhs_max_mode_indices = left_gate_str[0]
+    if not pure:
+        if mode1 == min_mode:
+            lhs_min_mode_indices += right_gate_str[0]
+            lhs_max_mode_indices += right_gate_str[2]
+            rhs_min_mode_indices += right_gate_str[1]
+            rhs_max_mode_indices += right_gate_str[3]
         else:
-            other_modes_indices = indices[
-                batch_offset
-                + 4 * mode_size : batch_offset
-                + 4 * mode_size
-                + mode_size * (num_modes - 2)
-            ]
-            # build equation
-            if mode1 == min_mode:
-                lhs_min_mode_indices = left_gate_str[1]
-                lhs_max_mode_indices = left_gate_str[3]
-                rhs_min_mode_indices = left_gate_str[0]
-                rhs_max_mode_indices = left_gate_str[2]
-            else:
-                lhs_min_mode_indices = left_gate_str[3]
-                lhs_max_mode_indices = left_gate_str[1]
-                rhs_min_mode_indices = left_gate_str[2]
-                rhs_max_mode_indices = left_gate_str[0]
-            if not pure:
-                if mode1 == min_mode:
-                    lhs_min_mode_indices += right_gate_str[0]
-                    lhs_max_mode_indices += right_gate_str[2]
-                    rhs_min_mode_indices += right_gate_str[1]
-                    rhs_max_mode_indices += right_gate_str[3]
-                else:
-                    lhs_min_mode_indices += right_gate_str[2]
-                    lhs_max_mode_indices += right_gate_str[0]
-                    rhs_min_mode_indices += right_gate_str[3]
-                    rhs_max_mode_indices += right_gate_str[1]
-            eqn_lhs = "{},{}{}{}{}{}{}".format(
-                batch_index + left_gate_str,
-                batch_index,
-                other_modes_indices[: min_mode * mode_size],
-                lhs_min_mode_indices,
-                other_modes_indices[min_mode * mode_size : (max_mode - 1) * mode_size],
-                lhs_max_mode_indices,
-                other_modes_indices[(max_mode - 1) * mode_size :],
-            )
-            if not pure:
-                eqn_lhs += "," + batch_index + right_gate_str
-            eqn_rhs = "".join(
-                [
-                    batch_index,
-                    other_modes_indices[: min_mode * mode_size],
-                    rhs_min_mode_indices,
-                    other_modes_indices[min_mode * mode_size : (max_mode - 1) * mode_size],
-                    rhs_max_mode_indices,
-                    other_modes_indices[(max_mode - 1) * mode_size :],
-                ]
-            )
-            eqn = eqn_lhs + "->" + eqn_rhs
-            einsum_inputs = [matrix, in_modes]
-            if not pure:
-                if batched:
-                    transpose_list = [0, 2, 1, 4, 3]
-                else:
-                    transpose_list = [1, 0, 3, 2]
-                einsum_inputs.append(tf.math.conj(tf.transpose(matrix, transpose_list)))
-            output = tf.einsum(eqn, *einsum_inputs)
-            return output
+            lhs_min_mode_indices += right_gate_str[2]
+            lhs_max_mode_indices += right_gate_str[0]
+            rhs_min_mode_indices += right_gate_str[3]
+            rhs_max_mode_indices += right_gate_str[1]
+    eqn_lhs = "{},{}{}{}{}{}{}".format(
+        batch_index + left_gate_str,
+        batch_index,
+        other_modes_indices[: min_mode * mode_size],
+        lhs_min_mode_indices,
+        other_modes_indices[min_mode * mode_size : (max_mode - 1) * mode_size],
+        lhs_max_mode_indices,
+        other_modes_indices[(max_mode - 1) * mode_size :],
+    )
+    if not pure:
+        eqn_lhs += "," + batch_index + right_gate_str
+    eqn_rhs = "".join(
+        [
+            batch_index,
+            other_modes_indices[: min_mode * mode_size],
+            rhs_min_mode_indices,
+            other_modes_indices[min_mode * mode_size : (max_mode - 1) * mode_size],
+            rhs_max_mode_indices,
+            other_modes_indices[(max_mode - 1) * mode_size :],
+        ]
+    )
+    eqn = eqn_lhs + "->" + eqn_rhs
+    einsum_inputs = [matrix, in_modes]
+    if not pure:
+        if batched:
+            transpose_list = [0, 2, 1, 4, 3]
+        else:
+            transpose_list = [1, 0, 3, 2]
+        einsum_inputs.append(tf.math.conj(tf.transpose(matrix, transpose_list)))
+    output = tf.einsum(eqn, *einsum_inputs)
+    return output
 
 
 def single_mode_superop(superop, mode, in_modes, pure=True, batched=False):
@@ -702,37 +701,34 @@ def single_mode_superop(superop, mode, in_modes, pure=True, batched=False):
         raise NotImplementedError(
             "The max number of supported modes for this operation is currently {}".format(max_len)
         )
-    else:
-        if pure:
-            in_modes = mixed(in_modes, batched)
 
-        # create equation
-        batch_index = indices[:batch_offset]
-        superop_indices = indices[batch_offset : batch_offset + 4]
-        state_indices = indices[batch_offset + 4 : batch_offset + 4 + 2 * num_modes]
-        left_unchanged_indices = state_indices[: 2 * mode]
-        right_unchanged_indices = state_indices[2 * mode : 2 * (num_modes - 1)]
-        eqn_lhs = ",".join(
-            [
-                batch_index + superop_indices,
-                batch_index
-                + left_unchanged_indices
-                + superop_indices[1:3]
-                + right_unchanged_indices,
-            ]
-        )
-        eqn_rhs = "".join(
-            [
-                batch_index,
-                left_unchanged_indices
-                + superop_indices[0]
-                + superop_indices[3]
-                + right_unchanged_indices,
-            ]
-        )
-        eqn = "->".join([eqn_lhs, eqn_rhs])
-        new_state = tf.einsum(eqn, superop, in_modes)
-        return new_state
+    if pure:
+        in_modes = mixed(in_modes, batched)
+
+    # create equation
+    batch_index = indices[:batch_offset]
+    superop_indices = indices[batch_offset : batch_offset + 4]
+    state_indices = indices[batch_offset + 4 : batch_offset + 4 + 2 * num_modes]
+    left_unchanged_indices = state_indices[: 2 * mode]
+    right_unchanged_indices = state_indices[2 * mode : 2 * (num_modes - 1)]
+    eqn_lhs = ",".join(
+        [
+            batch_index + superop_indices,
+            batch_index + left_unchanged_indices + superop_indices[1:3] + right_unchanged_indices,
+        ]
+    )
+    eqn_rhs = "".join(
+        [
+            batch_index,
+            left_unchanged_indices
+            + superop_indices[0]
+            + superop_indices[3]
+            + right_unchanged_indices,
+        ]
+    )
+    eqn = "->".join([eqn_lhs, eqn_rhs])
+    new_state = tf.einsum(eqn, superop, in_modes)
+    return new_state
 
 
 ###################################################################
@@ -1066,7 +1062,7 @@ def partial_trace(system, mode, state_is_pure, batched=False):
     # tensorflow trace implementation
     # requires subsystem to be traced out to be at end
     # i.e., ab...klmnop...yz goes to ab...klop...yzmn
-    indices_list = [m for m in range(batch_offset + 2 * num_modes)]
+    indices_list = list(range(batch_offset + 2 * num_modes))
     dim_list = (
         indices_list[: batch_offset + 2 * mode]
         + indices_list[batch_offset + 2 * (mode + 1) :]
@@ -1077,11 +1073,14 @@ def partial_trace(system, mode, state_is_pure, batched=False):
     return reduced_state
 
 
-def reduced_density_matrix(system, mode, state_is_pure, batched=False):
+def reduced_density_matrix(system, modes, state_is_pure, batched=False):
     """
-    Trace out all subsystems except 'mode' from 'system'.
+    Trace out all subsystems except those specified in ``modes`` from ``system``. ``modes`` can be either an int or a list.
     This operation always returns a mixed state, since we do not know in advance if a mode is entangled with others.
     """
+    if isinstance(modes, int):
+        modes = [modes]
+
     if state_is_pure:
         reduced_state = mixed(system, batched)
     else:
@@ -1092,9 +1091,11 @@ def reduced_density_matrix(system, mode, state_is_pure, batched=False):
     else:
         batch_offset = 0
     num_modes = (num_indices - batch_offset) // 2  # always mixed
+    removed_cnt = 0
     for m in range(num_modes):
-        if m != mode:
-            reduced_state = partial_trace(reduced_state, m, False, batched)
+        if m not in modes:
+            reduced_state = partial_trace(reduced_state, m - removed_cnt, False, batched)
+            removed_cnt += 1
     return reduced_state
 
 
diff --git a/strawberryfields/backends/tfbackend/states.py b/strawberryfields/backends/tfbackend/states.py
--- a/strawberryfields/backends/tfbackend/states.py
+++ b/strawberryfields/backends/tfbackend/states.py
@@ -95,7 +95,7 @@ def fock_prob(self, n, **kwargs):
         if len(n) != self.num_modes:
             raise ValueError("List length should be equal to number of modes")
 
-        elif max(n) >= self.cutoff_dim:
+        if max(n) >= self.cutoff_dim:
             raise ValueError("Can't get distribution beyond truncation level")
 
         state = self.data
@@ -288,7 +288,7 @@ def fidelity_vacuum(self, **kwargs):
 
         return v
 
-    def is_vacuum(self, tol=0.0, **kwargs):
+    def is_vacuum(self, tol=0.0, **kwargs):  # pylint:disable=unused-argument
         r"""
         Computes a boolean which indicates whether the state is the vacuum state. May be numerical or symbolic.
 
@@ -327,8 +327,7 @@ def reduced_dm(self, modes, **kwargs):
             )
 
         reduced = self.dm(**kwargs)
-        for m in modes:
-            reduced = reduced_density_matrix(reduced, m, False, batched=self.batched)
+        reduced = reduced_density_matrix(reduced, modes, False, batched=self.batched)
 
         s = tf.identity(reduced, name="reduced_density_matrix")
         return s
@@ -547,11 +546,11 @@ def poly_quad_expectation(self, A, d=None, k=0, phi=0, **kwargs):  # pragma: no
 
         if not self.batched:
             return super().poly_quad_expectation(A, d, k, phi, **kwargs)
-        else:
-            raise NotImplementedError(
-                "Calculation of multi-mode quadratic expectation values is currently "
-                "only supported when batched=False."
-            )
+
+        raise NotImplementedError(
+            "Calculation of multi-mode quadratic expectation values is currently "
+            "only supported when batched=False."
+        )
 
     @property
     def batched(self):


Test patch:
diff --git a/tests/backend/test_tf_ops.py b/tests/backend/test_tf_ops.py
new file mode 100644
--- /dev/null
+++ b/tests/backend/test_tf_ops.py
@@ -0,0 +1,78 @@
+# Copyright 2020 Xanadu Quantum Technologies Inc.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+r""" Tests for the file tfbackend/ops.py"""
+
+import pytest
+
+import numpy as np
+tf = pytest.importorskip("tensorflow", minversion="2.0")
+
+from strawberryfields.backends.tfbackend.ops import reduced_density_matrix
+
+@pytest.mark.backends("tf")
+class TestTFOps:
+    """Testing for tfbackend/ops.py"""
+
+    @pytest.mark.parametrize("num_modes", [2, 3, 4])
+    def test_reduced_density_matrix_fock_states(self, num_modes, setup_backend, cutoff, tol):
+        """Test the reduced_density_matrix returns the correct reduced density matrices
+        when prepared with initial fock states"""
+
+        zero_photon_state = np.zeros([cutoff])
+        zero_photon_state[0] = 1.
+        one_photon_state = np.zeros([cutoff])
+        one_photon_state[1] = 1.
+
+        # create a single-photon state in the second mode
+        state = np.outer(zero_photon_state, one_photon_state)
+        for _ in range(2, num_modes):
+            state = np.multiply.outer(state, zero_photon_state)
+        state = tf.constant(state)
+
+        # get reduced density matrix of last subsystem
+        mode = num_modes - 1
+        reduced_dm = reduced_density_matrix(state, mode, state_is_pure=True)
+
+        if num_modes == 2:
+            expected = np.multiply.outer(one_photon_state, one_photon_state)
+        else:
+            expected = np.multiply.outer(zero_photon_state, zero_photon_state)
+        expected = tf.constant(expected)
+
+        assert np.allclose(reduced_dm, expected, atol=tol, rtol=0)
+
+    @pytest.mark.parametrize("modes,einstr", [
+        ([0], "abccee->ab"),
+        ([1], "aacdee->cd"),
+        ([2], "aaccef->ef"),
+        ([0, 1], "abcdee->abcd"),
+        ([0, 2], "abccef->abef"),
+        ([1, 2], "aacdef->cdef"),
+        ([0, 1, 2], "abcdef->abcdef")
+    ])
+    def test_reduced_density_matrix_multiple_modes(self, setup_backend, cutoff, modes, einstr, tol):
+        """Test that reduced_density_matrix returns the correct reduced density matrices."""
+        state = np.zeros([cutoff, cutoff, cutoff])
+        state[0, 0, 0] = 1 / np.sqrt(2)
+        state[1, 1, 1] = 1 / np.sqrt(2)
+        state = np.multiply.outer(state, state.conj())
+
+        # transpose the state due to different conventions
+        state = state.transpose([0, 3, 1, 4, 2, 5])
+        state = tf.constant(state)
+
+        reduced_dm = reduced_density_matrix(state, modes, state_is_pure=False)
+        expected = tf.constant(np.einsum(einstr, state))
+
+        assert np.allclose(reduced_dm, expected, atol=tol, rtol=0)
diff --git a/tests/integration/test_tf_integration.py b/tests/integration/test_tf_integration.py
--- a/tests/integration/test_tf_integration.py
+++ b/tests/integration/test_tf_integration.py
@@ -22,8 +22,7 @@
 
 tf = pytest.importorskip("tensorflow", minversion="2.0")
 
-from strawberryfields.ops import Dgate, Sgate, MeasureX, Thermal, BSgate, Ket, S2gate
-import strawberryfields.parameters
+from strawberryfields.ops import Dgate, Sgate, MeasureX, Thermal, BSgate, Ket, S2gate, Fock
 
 
 # this test file is only supported by the TF backend
@@ -249,7 +248,7 @@ def test_eng_run_state_mean_photon(self, setup_eng, tol):
         assert np.allclose(var, ref_var, atol=tol, rtol=0.0)
 
 
-class TestTwoModeSymbolic:
+class TestTwoModes:
     """Tests for workflows on two modes."""
 
     def test_eval_true_state_all_fock_probs(self, setup_eng, cutoff, batch_size, tol):
@@ -301,6 +300,32 @@ def test_eval_true_state_fock_prob(self, setup_eng, cutoff, tol):
         assert np.allclose(prob, ref_prob, atol=tol, rtol=0.0)
 
 
+class TestThreeModes:
+    """Tests for workflows on three modes."""
+
+    def test_three_mode_fock_state_mean_photon(self, setup_eng, tol):
+        """Tests that the mean photon numbers of a three-mode system,
+        where one mode is prepared in a Fock state, are correct."""
+
+        eng, prog = setup_eng(3)
+
+        with prog.context as q:
+            Fock(1) | q[1]
+
+        result = eng.run(prog)
+        state = result.state
+        nbar0, var_n0 = state.mean_photon(0)
+        nbar1, var_n1 = state.mean_photon(1)
+        nbar2, var_n2 = state.mean_photon(2)
+
+        assert np.allclose(nbar0, 0., atol=tol, rtol=0)
+        assert np.allclose(nbar1, 1., atol=tol, rtol=0)
+        assert np.allclose(nbar2, 0., atol=tol, rtol=0)
+
+        assert np.allclose(var_n0, 0., atol=tol, rtol=0)
+        assert np.allclose(var_n1, 0., atol=tol, rtol=0)
+        assert np.allclose(var_n2, 0., atol=tol, rtol=0)
+
 class TestGradient:
     """Integration tests for the gradient computation"""