reduced_dm not correct for subsystems of more than 1 mode

[co9olguy]

Issue description

• Expected behavior: Calling reduced_dm with a list of modes with length longer than one should give an output with the correct shape (e.g., for a system with three modes, using modes=[0,1] should return a density matrix with shape (D, D, D, D), where D is the cutoff dimension)

• Actual behavior: The output tensors do not have the correct shape, they are missing required axes, or become scalars. I have diagnosed the bug to the following lines in tfbackend/states.py:

  for m in modes:
          reduced = reduced_density_matrix(reduced, m, False, batched=self.batched)

  The function reduced_density_matrix, according to its docstring, traces out all modes except m. Thus, the loop above always produces a density matrix of shape (D, D) after its first iteration. This is not how one produces reduced density matrices when more than two subsystems are kept

• Reproduces how often: occurs when using the TF backend and more than 2 modes, but not the Fock backend

• System information:

  sf.about()
  Python 3.8.2 (default, May  7 2020, 20:00:49) 
  Type 'copyright', 'credits' or 'license' for more information
  IPython 7.18.1 -- An enhanced Interactive Python. Type '?' for help.
  PyDev console: using IPython 7.18.1
  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.4.0-48-generic-x86_64-with-glibc2.10
  Installation path:         /home/nathan/dev/strawberryfields/strawberryfields
  Strawberry Fields version: 0.15.0-dev0
  Numpy version:             1.18.4
  Scipy version:             1.5.2
  SymPy version:             1.6.2
  NetworkX version:          2.5
  The Walrus version:        0.13.0
  Blackbird version:         0.2.4
  TensorFlow version:        2.3.1

Source code and tracebacks

Minimal (non)working example:

import strawberryfields as sf

modes = 3
fock_cutoff = 3
prog = sf.Program(modes)
eng = sf.Engine("tf", backend_options={"cutoff_dim": fock_cutoff})
result = eng.run(prog)
state = result.state
d01 = state.reduced_dm([0,1])
d02 = state.reduced_dm([0,2])
print(d01.shape)
print(d02.shape)
>>> ()  # should be (3,3,3,3)
>>> ()  # should be (3,3,3,3)

[nquesada]

Was this closed here https://github.com/XanaduAI/strawberryfields/pull/471 @thisac ?

[thisac]

@nquesada Yes, this should've been closed by #471. Thanks for spotting it!