# Source code for thewalrus.quantum.fock_tensors

# Copyright 2019-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

# 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.
"""
Set of functions for calculating various state representations, probabilities and
classical subsystems of Gaussian states.
"""
# pylint: disable=too-many-arguments

from itertools import count, product, chain

from collections import OrderedDict

import numpy as np

from scipy.special import factorial as fac
from numba import jit

from ..symplectic import expand, is_symplectic, reduced_state

from .._hafnian import hafnian, hafnian_repeated, reduction
from .._hermite_multidimensional import hermite_multidimensional, hafnian_batched, interferometer

from .conversions import (
Amat,
Qmat,
reduced_gaussian,
complex_to_real_displacements,
)

from .gaussian_checks import is_classical_cov, is_pure_cov, is_valid_cov

def pure_state_amplitude(mu, cov, i, include_prefactor=True, tol=1e-10, hbar=2, check_purity=True):
r"""Returns the :math:\langle i | \psi\rangle element of the state ket
of a Gaussian state defined by covariance matrix cov.

Args:
mu (array): length-:math:2N quadrature displacement vector
cov (array): length-:math:2N covariance matrix
i (list): list of amplitude elements
include_prefactor (bool): if True, the prefactor is automatically calculated
used to scale the result.
tol (float): tolerance for determining if displacement is negligible
hbar (float): the value of :math:\hbar in the commutation
relation :math:[\x,\p]=i\hbar.
check_purity (bool): if True, the purity of the Gaussian state is checked
before calculating the state vector.

Returns:
complex: the pure state amplitude
"""
if check_purity:
if not is_pure_cov(cov, hbar=hbar, rtol=1e-05, atol=1e-08):
raise ValueError("The covariance matrix does not correspond to a pure state")

rpt = i
beta = complex_to_real_displacements(mu, hbar=hbar)
Q = Qmat(cov, hbar=hbar)
A = Amat(cov, hbar=hbar)
(n, _) = cov.shape
N = n // 2
B = A[0:N, 0:N].conj()
alpha = beta[0:N]

if np.linalg.norm(alpha) < tol:
# no displacement
if np.prod([k + 1 for k in rpt]) ** (1 / len(rpt)) < 3:
B_rpt = reduction(B, rpt)
haf = hafnian(B_rpt)
else:
haf = hafnian_repeated(B, rpt)
else:
gamma = alpha - B @ np.conj(alpha)
if np.prod([k + 1 for k in rpt]) ** (1 / len(rpt)) < 3:
B_rpt = reduction(B, rpt)
np.fill_diagonal(B_rpt, reduction(gamma, rpt))
haf = hafnian(B_rpt, loop=True)
else:
haf = hafnian_repeated(B, rpt, mu=gamma, loop=True)

if include_prefactor:
pref = np.exp(-0.5 * (np.linalg.norm(alpha) ** 2 - alpha.conj() @ B @ alpha.conj()))
haf *= pref

return haf / np.sqrt(np.prod(fac(rpt)) * np.sqrt(np.linalg.det(Q)))

def state_vector(
mu, cov, post_select=None, normalize=False, cutoff=5, hbar=2, check_purity=True, **kwargs
):
r"""Returns the state vector of a (PNR post-selected) Gaussian state.

The resulting density matrix will have shape

.. math:: \underbrace{D\times D \times \cdots \times D}_M

where :math:D is the Fock space cutoff, and :math:M is the
number of *non* post-selected modes, i.e. M = len(mu)//2 - len(post_select).

If post_select is None then the density matrix elements are calculated using
the multidimensional Hermite polynomials which provide a significantly faster
evaluation.

Args:
mu (array): length-:math:2N means vector in xp-ordering
cov (array): :math:2N\times 2N covariance matrix in xp-ordering
post_select (dict): dictionary containing the post-selected modes, of
the form {mode: value}.
normalize (bool): If True, a post-selected density matrix is re-normalized.
cutoff (dim): the final length (i.e., Hilbert space dimension) of each
mode in the density matrix.
hbar (float): the value of :math:\hbar in the commutation
relation :math:[\x,\p]=i\hbar.
check_purity (bool): if True, the purity of the Gaussian state is checked
before calculating the state vector.

Keyword Args:
choi_r (float or None): Value of the two-mode squeezing parameter used in Choi-Jamiolkoski
trick in :func:~.fock_tensor. This keyword argument should only be used when state_vector
is called by :func:~.fock_tensor.

Returns:
np.array[complex]: the state vector of the Gaussian state
"""
if check_purity:
if not is_pure_cov(cov, hbar=hbar, rtol=1e-05, atol=1e-08):
raise ValueError("The covariance matrix does not correspond to a pure state")

beta = complex_to_real_displacements(mu, hbar=hbar)
A = Amat(cov, hbar=hbar)
Q = Qmat(cov, hbar=hbar)

(n, _) = cov.shape
N = n // 2

B = A[0:N, 0:N]
alpha = beta[0:N]
gamma = np.conj(alpha) - B @ alpha
prefexp = -0.5 * (np.linalg.norm(alpha) ** 2 - alpha @ B @ alpha)
pref = np.exp(prefexp.conj())
if post_select is None:
choi_r = kwargs.get("choi_r", None)
if choi_r is None:
denom = np.sqrt(np.sqrt(np.linalg.det(Q).real))
else:
rescaling = np.concatenate(
[np.ones([N // 2]), (1.0 / np.tanh(choi_r)) * np.ones([N // 2])]
)
B = np.diag(rescaling) @ B @ np.diag(rescaling)
gamma = rescaling * gamma
denom = np.sqrt(np.sqrt(np.linalg.det(Q / np.cosh(choi_r)).real))

psi = pref * hafnian_batched(B.conj(), cutoff, mu=gamma.conj(), renorm=True) / denom
else:
M = N - len(post_select)
psi = np.zeros([cutoff] * (M), dtype=np.complex128)

for idx in product(range(cutoff), repeat=M):
el = []

counter = count(0)
modes = (np.arange(N)).tolist()
el = [post_select[i] if i in post_select else idx[next(counter)] for i in modes]
psi[idx] = pure_state_amplitude(
mu, cov, el, check_purity=False, include_prefactor=False, hbar=hbar
)

psi = psi * pref

if normalize:
norm = np.sqrt(np.sum(np.abs(psi) ** 2))
psi = psi / norm

return psi

def density_matrix_element(mu, cov, i, j, include_prefactor=True, tol=1e-10, hbar=2):
r"""Returns the :math:\langle i | \rho | j \rangle element of the density matrix
of a Gaussian state defined by covariance matrix cov.

Args:
mu (array): length-:math:2N quadrature displacement vector
cov (array): length-:math:2N covariance matrix
i (list): list of density matrix rows
j (list): list of density matrix columns
include_prefactor (bool): if True, the prefactor is automatically calculated
used to scale the result.
tol (float): tolerance for determining if displacement is negligible
hbar (float): the value of :math:\hbar in the commutation
relation :math:[\x,\p]=i\hbar.

Returns:
complex: the density matrix element
"""
rpt = i + j
beta = complex_to_real_displacements(mu, hbar=hbar)
A = Amat(cov, hbar=hbar)
if np.linalg.norm(beta) < tol:
# no displacement
if np.prod([k + 1 for k in rpt]) ** (1 / len(rpt)) < 3:
A_rpt = reduction(A, rpt)
haf = hafnian(A_rpt)
else:
haf = hafnian_repeated(A, rpt)
else:
# replace the diagonal of A with gamma
gamma = beta.conj() - A @ beta
if np.prod([k + 1 for k in rpt]) ** (1 / len(rpt)) < 3:
A_rpt = reduction(A, rpt)
np.fill_diagonal(A_rpt, reduction(gamma, rpt))
haf = hafnian(A_rpt, loop=True)
else:
haf = hafnian_repeated(A, rpt, mu=gamma, loop=True)

if include_prefactor:
haf *= _prefactor(mu, cov, hbar=hbar)

return haf / np.sqrt(np.prod(fac(rpt)))

def density_matrix(mu, cov, post_select=None, normalize=False, cutoff=5, hbar=2):
r"""Returns the density matrix of a (PNR post-selected) Gaussian state.

The resulting density matrix will have shape

.. math:: \underbrace{D\times D \times \cdots \times D}_{2M}

where :math:D is the Fock space cutoff, and :math:M is the
number of *non* post-selected modes, i.e. M = len(mu)//2 - len(post_select).

Note that we use the Strawberry Fields convention for indexing the density
matrix; the first two dimensions correspond to subsystem 1, the second two
dimensions correspond to subsystem 2, etc.
If post_select is None then the density matrix elements are calculated using
the multidimensional Hermite polynomials which provide a significantly faster
evaluation.

Args:
mu (array): length-:math:2N means vector in xp-ordering
cov (array): :math:2N\times 2N covariance matrix in xp-ordering
post_select (dict): dictionary containing the post-selected modes, of
the form {mode: value}. If post_select is None the whole non post-selected density matrix
is calculated directly using (multidimensional) Hermite polynomials, which is significantly faster
than calculating one hafnian at a time.
normalize (bool): If True, a post-selected density matrix is re-normalized.
cutoff (dim): the final length (i.e., Hilbert space dimension) of each
mode in the density matrix.
hbar (float): the value of :math:\hbar in the commutation
relation :math:[\x,\p]=i\hbar.

Returns:
np.array[complex]: the density matrix of the Gaussian state
"""
N = len(mu) // 2
pref = _prefactor(mu, cov, hbar=hbar)

if post_select is None:
A = Amat(cov, hbar=hbar).conj()
sf_order = tuple(chain.from_iterable([[i, i + N] for i in range(N)]))

if np.allclose(mu, np.zeros_like(mu)):
tensor = pref * hermite_multidimensional(-A, cutoff, renorm=True, modified=True)
return tensor.transpose(sf_order)
beta = complex_to_real_displacements(mu, hbar=hbar)
y = beta - A @ beta.conj()
tensor = pref * hermite_multidimensional(-A, cutoff, y=y, renorm=True, modified=True)
return tensor.transpose(sf_order)

M = N - len(post_select)
rho = np.zeros([cutoff] * (2 * M), dtype=np.complex128)

for idx in product(range(cutoff), repeat=2 * M):
el = []

counter = count(0)
modes = (np.arange(2 * N) % N).tolist()
el = [post_select[i] if i in post_select else idx[next(counter)] for i in modes]

el = np.array(el).reshape(2, -1)
el0 = el[0].tolist()
el1 = el[1].tolist()

sf_idx = np.array(idx).reshape(2, -1)
sf_el = tuple(sf_idx[::-1].T.flatten())

rho[sf_el] = density_matrix_element(mu, cov, el0, el1, include_prefactor=False, hbar=hbar)

rho *= pref

if normalize:
# construct the standard 2D density matrix, and take the trace
new_ax = np.arange(2 * M).reshape([M, 2]).T.flatten()
tr = np.trace(rho.transpose(new_ax).reshape([cutoff**M, cutoff**M])).real
# renormalize
rho /= tr

return rho

[docs]
def fock_tensor(
S,
alpha,
cutoff,
choi_r=np.arcsinh(1.0),
check_symplectic=True,
sf_order=False,
rtol=1e-05,
atol=1e-08,
):
r"""
Calculates the Fock representation of a Gaussian unitary parametrized by
the symplectic matrix S and the displacements alpha up to cutoff in Fock space.

Args:
S (array): symplectic matrix
alpha (array): complex vector of displacements
cutoff (int): cutoff in Fock space
choi_r (float): squeezing parameter used for the Choi expansion
check_symplectic (boolean): checks whether the input matrix is symplectic
sf_order (boolean): reshapes the tensor so that it follows the sf ordering of indices
rtol (float): the relative tolerance parameter used in np.allclose
atol (float): the absolute tolerance parameter used in np.allclose

Return:
(array): Tensor containing the Fock representation of the Gaussian unitary
"""
# Check the matrix is symplectic
if check_symplectic:
if not is_symplectic(S, rtol=rtol, atol=atol):
raise ValueError("The matrix S is not symplectic")

# And that S and alpha have compatible dimensions
m, _ = S.shape
l = m // 2
if l != len(alpha):
raise ValueError("The matrix S and the vector alpha do not have compatible dimensions")
# Check if S corresponds to an interferometer, if so use optimized routines
if np.allclose(S @ S.T, np.identity(m), rtol=rtol, atol=atol) and np.allclose(
alpha, 0, rtol=rtol, atol=atol
):
reU = S[:l, :l]
imU = S[:l, l:]
U = reU - 1j * imU
Ub = np.block([[0 * U, -U], [-U.T, 0 * U]])
tensor = interferometer(Ub, cutoff)
else:
# Construct the covariance matrix of l two-mode squeezed vacua pairing modes i and i+l
ch = np.cosh(choi_r) * np.identity(l)
sh = np.sinh(choi_r) * np.identity(l)
zh = np.zeros([l, l])
Schoi = np.block([[ch, sh, zh, zh], [sh, ch, zh, zh], [zh, zh, ch, -sh], [zh, zh, -sh, ch]])
# And then its Choi expanded symplectic
S_exp = expand(S, list(range(l)), 2 * l) @ Schoi
# And this is the corresponding covariance matrix
cov = S_exp @ S_exp.T
alphat = np.array(list(alpha) + ([0] * l))
x = 2 * alphat.real
p = 2 * alphat.imag
mu = np.concatenate([x, p])

tensor = state_vector(
mu,
cov,
normalize=False,
cutoff=cutoff,
hbar=2,
check_purity=False,
choi_r=choi_r,
)

if sf_order:
sf_indexing = tuple(chain.from_iterable([[i, i + l] for i in range(l)]))
return tensor.transpose(sf_indexing)

return tensor

[docs]
def probabilities(mu, cov, cutoff, parallel=False, hbar=2.0, rtol=1e-05, atol=1e-08):
r"""Generate the Fock space probabilities of a Gaussian state up to a Fock space cutoff.

Args:
mu (array): vector of means of length 2*n_modes
cov (array): covariance matrix of shape [2*n_modes, 2*n_modes]
cutoff (int): cutoff in Fock space
parallel (bool): if True, uses dask for parallelization instead of OpenMP
hbar (float): value of :math:\hbar in the commutation relation :math;[\hat{x}, \hat{p}]=i\hbar
rtol (float): the relative tolerance parameter used in np.allclose
atol (float): the absolute tolerance parameter used in np.allclose

Returns:
(array): Fock space probabilities up to cutoff. The shape of this tensor is [cutoff]*num_modes.
"""
if is_pure_cov(
cov, hbar=hbar, rtol=rtol, atol=atol
):  # Check if the covariance matrix cov is pure
return np.abs(state_vector(mu, cov, cutoff=cutoff, hbar=hbar, check_purity=False)) ** 2
num_modes = len(mu) // 2

if parallel:
compute_list = []
# create a list of parallelizable computations
for i in product(range(cutoff), repeat=num_modes):
compute_list.append(dask.delayed(density_matrix_element)(mu, cov, i, i, hbar=hbar))

probs = np.maximum(
).reshape([cutoff] * num_modes)
# maximum is needed because sometimes a probability is very close to zero from below
else:
probs = np.zeros([cutoff] * num_modes)
for i in product(range(cutoff), repeat=num_modes):
probs[i] = np.maximum(
0.0, np.real_if_close(density_matrix_element(mu, cov, i, i, hbar=hbar))
)
# maximum is needed because sometimes a probability is very close to zero from below
return probs

[docs]
@jit(nopython=True)
def loss_mat(eta, cutoff):  # pragma: no cover
r"""Constructs a binomial loss matrix with transmission eta up to n photons.

Args:
eta (float): Transmission coefficient. eta=0.0 corresponds to complete loss and eta=1.0 corresponds to no loss.
cutoff (int): cutoff in Fock space.

Returns:
array: :math:n\times n matrix representing the loss.
"""
# If full transmission return the identity

if eta < 0.0 or eta > 1.0:
raise ValueError("The transmission parameter eta should be a number between 0 and 1.")

if eta == 1.0:
return np.identity(cutoff)

# Otherwise construct the matrix elements recursively
lm = np.zeros((cutoff, cutoff))
mu = 1.0 - eta
lm[:, 0] = mu ** (np.arange(cutoff))
for i in range(cutoff):
for j in range(1, i + 1):
lm[i, j] = lm[i, j - 1] * (eta / mu) * (i - j + 1) / (j)
return lm

[docs]
def update_probabilities_with_loss(etas, probs):
"""Given a list of transmissivities a tensor of probabilitites, calculate
an updated tensor of probabilities after loss is applied.

Args:
etas (list): List of transmissitivities describing the loss in each of the modes
probs (array): Array of probabilitites in the different modes

Returns:
array: List of loss-updated probabilities with the same shape as probs.
"""

probs_shape = probs.shape
if len(probs_shape) != len(etas):
raise ValueError(
"The list of transmission etas and the tensor of probabilities probs have incompatible dimensions."
)

alphabet = "abcdefghijklmnopqrstuvwxyz"
cutoff = probs_shape[0]
for i, eta in enumerate(etas):
einstrings = "ij,{}i...->{}j...".format(alphabet[:i], alphabet[:i])

qein = np.zeros_like(probs)
qein = np.einsum(einstrings, loss_mat(eta, cutoff), probs)
probs = np.copy(qein)
return qein

@jit(nopython=True)
def _update_1d(probs, one_d, cutoff):  # pragma: no cover
"""Performs a convolution of the two arrays. The first one does not need to be one dimensional, which is why we do not use np.convolve.

Args:
probs (array): (multidimensional) array
one_d (array): one dimensional array
cutoff (int): cutoff in Fock space for the first array

Returns:
(array): the convolution of the two arrays, with the same shape as probs.
"""
new_d = np.zeros_like(probs)
for i in range(cutoff):
for j in range(min(i + 1, len(one_d))):
new_d[i] += probs[i - j] * one_d[j]
return new_d

[docs]
def update_probabilities_with_noise(probs_noise, probs):
"""Given a list of noise probability distributions for each of the modes and a tensor of
probabilitites, calculate an updated tensor of probabilities after noise is applied.

Args:
probs_noise (list): List of probability distributions describing the noise in each of the modes
probs (array): Array of probabilitites in the different modes

Returns:
array: List of noise-updated probabilities with the same shape as probs.
"""
probs_shape = probs.shape
num_modes = len(probs_shape)
cutoff = probs_shape[0]
if num_modes != len(probs_noise):
raise ValueError(
"The list of probability distributions probs_noise and the tensor of probabilities probs have incompatible dimensions."
)

for k in range(num_modes):  # update one mode at a time
perm = np.arange(num_modes)
perm[0] = k
perm[k] = 0
one_d = probs_noise[k]
probs_masked = np.transpose(probs, axes=perm)
probs = np.transpose(probs_masked, axes=perm)
return probs

[docs]
def find_classical_subsystem(cov, hbar=2, atol=1e-08):
"""Find the largest integer k so that subsystem in modes [0,1,...,k-1] is a classical state.

Args:
cov (array): a covariance matrix
hbar (float): value of hbar in the uncertainty relation
atol (float): the absolute tolerance parameter used when determining if the state is classical

Returns:
int: the largest k so that modes [0,1,...,k-1] are in a classical state.
"""
n, _ = cov.shape
nmodes = n // 2
if is_classical_cov(cov, hbar=hbar, atol=atol):
return nmodes
k = 0
mu = np.zeros(n)
is_classical = True
while is_classical:
_, Vk = reduced_gaussian(mu, cov, list(range(k + 1)))
is_classical = is_classical_cov(Vk, hbar=hbar, atol=atol)
k += 1
return k - 1

def _prefactor(mu, cov, hbar=2):
r"""Returns the prefactor.

.. math:: prefactor = \frac{e^{-\beta Q^{-1}\beta^*/2}}{n_1!\cdots n_m! \sqrt{|Q|}}

Args:
mu (array): length-:math:2N vector of mean values :math:[\alpha,\alpha^*]
cov (array): length-:math:2N xp-covariance matrix

Returns:
float: the prefactor
"""
Q = Qmat(cov, hbar=hbar)
beta = complex_to_real_displacements(mu, hbar=hbar)
Qinv = np.linalg.inv(Q)
return np.exp(-0.5 * beta @ Qinv @ beta.conj()) / np.sqrt(np.linalg.det(Q))

[docs]
def tvd_cutoff_bounds(mu, cov, cutoff, hbar=2, check_is_valid_cov=True, rtol=1e-05, atol=1e-08):
r"""Gives bounds of the total variation distance between the exact Gaussian Boson Sampling
distribution extending to infinity in Fock space and the ones truncated by any value between 0
and the user provided cutoff.

For the derivation see Appendix B of 'Exact simulation of Gaussian boson sampling in polynomial space and exponential time',
Quesada and Arrazola <https://journals.aps.org/prresearch/abstract/10.1103/PhysRevResearch.2.023005>_.

Args:
mu (array): vector of means of the Gaussian state
cov (array): covariance matrix of the Gaussian state
cutoff (int): cutoff in Fock space
check_is_valid_cov (bool): verify that the covariance matrix is physical
hbar (float): value of hbar in the uncertainty relation
rtol (float): the relative tolerance parameter used in np.allclose
atol (float): the absolute tolerance parameter used in np.allclose

Returns:
(array): values of the bound for different local Fock space dimensions up to cutoff
"""
if check_is_valid_cov:
if not is_valid_cov(cov, hbar=hbar, rtol=rtol, atol=atol):
raise ValueError("The input covariance matrix violates the uncertainty relation.")
nmodes = cov.shape[0] // 2
bounds = np.zeros([cutoff])
for i in range(nmodes):
mu_red, cov_red = reduced_state(mu, cov, [i])
ps = np.real_if_close(np.diag(density_matrix(mu_red, cov_red, cutoff=cutoff, hbar=hbar)))
bounds += 1 - np.cumsum(ps)
return bounds

[docs]
def n_body_marginals(mean, cov, cutoff, n, hbar=2):
r"""Calculates the first n-body marginals of a Gaussian state.

For an M-mode Gaussian state there exists a photon number distribution with probability mass function
:math:p[i_0,i_1,\ldots, i_{M-1}]. The function n_body_marginals calculates the first n-body marginals
of the (all-mode) probability distribution :math:p. The :math:n=1 marginals or single body marginals
are simply the probability that mode :math:k has :math:i photons, i.e. :math:p_k[i].
For :math:n=2 one obtains the two-body probabilities. For two modes :math:k and :math:l this is a
two dimensional probability distribution :math:p_{k,l}[i,j].
Note that these marginals have interesting permutation properties, for example :math:p_{k,l}[i,j] = p_{l,k}[j,i].

The function provided here takes advantage of these symmetries to minimize the amount of calculations.
The return of this function is a list of tensors where the first entry contains the one-body marginals of the :math:M modes
(giving a tensor of shape (M, cutoff)), the second entry contains the two-body marginals (giving a tensor of shape (M,M,cutoff, cutoff))
and so on and so forth.

To be clever about not calculating things that can be obtained by permutations it checks whether the index vector representing the modes is sorted.
From the way itertools.product works we know that it will always produce a sorted index vector before generating any of its unordered permutations.
Thus whenever the index vector is ordered we perform the numerical calculation.

If it is an unsorted index vector it realizes, in the if statement, that it can be obtained by permuting the
marginal distribution of something that has already been calculated.

Args:
mean (array): length-:math:2N quadrature displacement vector
cov (array): length-:math:2N covariance matrix
cutoff (int): cutoff in Fock space
n (int): order of the correlations
hbar (float): the value of :math:\hbar in the commutation
relation :math:[\x,\p]=i\hbar.

Returns:
list[array]: List with arrays containing the :math:1,..,n body marginal
distributions of the modes
"""
M = len(mean)
if (M, M) != cov.shape:
raise ValueError("The covariance matrix and vector of means have incompatible dimensions")
if M % 2 != 0:
raise ValueError("The vector of means is not of even dimensions")
M = M // 2
if M < n:
raise ValueError("The order of the correlations is higher than the number of modes")

marginal = [np.zeros(([M] * i) + ([cutoff] * i)) for i in range(1, n + 1)]

for ind in product(range(M), repeat=n):
modes = list(set(ind))
acc = len(modes) - 1
if list(ind) == sorted(ind):
sub_mean, sub_cov = reduced_state(mean, cov, modes)  # this happens in phase space
marginal[acc][tuple(modes)] = probabilities(sub_mean, sub_cov, cutoff, hbar=hbar)
else:
modes_usrt = list(OrderedDict.fromkeys(ind))
perm = np.argsort(modes_usrt)
marginal[acc][tuple(modes_usrt)] = marginal[acc][tuple(modes)].transpose(perm)
return marginal