AttributeError when optimizing purity of reduced density matrix

[sillygoose]

Hello! I am trying to optimize the linear entropy of a time evolved, then traced out density matrix using pennylane. The Hamiltonian used to time evolve my state is what is being parameterized.

However, I receive an error: AttributeError: 'ArrayBox' object has no attribute 'dot'. Any help in solving this error and/or pointing me towards a better way to optimize the linear entropy of a time evolved, then traced out density matrix would be greatly appreciated! What I am doing is precisely the following.

First, I am dealing with qubit systems. Let us fix the system to be a system of three qubits and define an initial density matrix of this three qubit system. I will also define the parameters to be optimized.

import pennylane as qml
from pennylane import numpy as np
import jax
import qutip as qtp
from scipy.optimize import minimize

n = 3 # total number of qubits
d = 2**n # dimension of composite system
H = isingHam # native Hamiltonian
tau = 1 # characteristic time

rho = np.zeros((d, d), dtype=np.complex128)
val = 0.25
rho[0, 0] = val
rho[1, 1] = val
rho[2, 2] = val
rho[3, 3] = val
rho = np.array(rho) # initial density matrix

wireList = list(range(n))
thetas = np.array(np.random.randn((d**2-1), requires_grad = True)) # random initial parameters for optimization

I also start with a fixed Hamiltonian, say the 1D Ising model Hamiltonian with boundary conditions:

coeffs = [0.5]*3
ops = [qml.PauliZ(0) @ qml.PauliZ(1), qml.PauliZ(1) @ qml.PauliZ(2), qml.PauliZ(0) @ qml.PauliZ(2)]
isingHam = qml.Hamiltonian(coeffs, ops) # native Hamiltonian

Next, I time evolve the initial density matrix by a “scrambled Hamiltonian” and then trace out of the environment of the initial density matrix. By “scrambled Hamiltonian”, I mean the fixed Hamiltonian defined earlier conjugated by a special unitary operator, parameterized by the thetas.

devRho = qml.device("default.mixed", wires=3)
@qml.qnode(devRho)
def processRho(n, H, tau, thetas, rho):
    '''
    Time evolves rho using scrambled Hamiltonian by characteristic time, then traces out of environment.
    Inputs:
     - n: total number of qubits
     - H: native Hamiltonian to be scrambled
     - tau: characteristic time
     - thetas: parameters to scramble H
     - rho: initial density matrix
    
    Outputs:
     - rhoTau: time evolved, then traced out density matrix
    '''
    qml.QubitDensityMatrix(rho, wires=wireList)
    qml.adjoint(qml.SpecialUnitary(thetas, wires=wireList))
    qml.ApproxTimeEvolution(H, tau, 100)
    qml.SpecialUnitary(thetas, wires=wireList)
    return qml.density_matrix([0])

#print(qml.draw(circuit)(n, isingHam, tau, thetas, rho))

I initialize the reduced density matrix returned from processRho() and take its purity. Finally, I define my cost function, i.e. the linear entropy:

@qml.qnode(devRho)
def purity(thetas):
    '''
    Initializes a density matrix and takes its purity.
    '''
    qml.QubitDensityMatrix(processRho(n, H, tau, thetas, rho), wires=0)
    return qml.purity(0)

def cost(thetas):
    '''
    Computes linear entropy from purity. 
    '''
    return 1 - purity(thetas)

Just as a simple test, I thought to test optimization using:

minimize(purity, thetas, method='BFGS', jac=qml.grad(cost, argnum=0))

However, running all of the code (in particular, running the optimization) gives me the error message:

---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
<ipython-input-69-1d027bdc7de1> in <module>
----> 1 minimize(purity, thetas, method='BFGS', jac=qml.grad(cost, argnum=0))

~/opt/anaconda3/lib/python3.8/site-packages/scipy/optimize/_minimize.py in minimize(fun, x0, args, method, jac, hess, hessp, bounds, constraints, tol, callback, options)
    612         return _minimize_cg(fun, x0, args, jac, callback, **options)
    613     elif meth == 'bfgs':
--> 614         return _minimize_bfgs(fun, x0, args, jac, callback, **options)
    615     elif meth == 'newton-cg':
    616         return _minimize_newtoncg(fun, x0, args, jac, hess, hessp, callback,

~/opt/anaconda3/lib/python3.8/site-packages/scipy/optimize/optimize.py in _minimize_bfgs(fun, x0, args, jac, callback, gtol, norm, eps, maxiter, disp, return_all, finite_diff_rel_step, **unknown_options)
   1133         maxiter = len(x0) * 200
   1134 
-> 1135     sf = _prepare_scalar_function(fun, x0, jac, args=args, epsilon=eps,
   1136                                   finite_diff_rel_step=finite_diff_rel_step)
   1137 

~/opt/anaconda3/lib/python3.8/site-packages/scipy/optimize/optimize.py in _prepare_scalar_function(fun, x0, jac, args, bounds, epsilon, finite_diff_rel_step, hess)
    259     # ScalarFunction caches. Reuse of fun(x) during grad
    260     # calculation reduces overall function evaluations.
--> 261     sf = ScalarFunction(fun, x0, args, grad, hess,
    262                         finite_diff_rel_step, bounds, epsilon=epsilon)
    263 

~/opt/anaconda3/lib/python3.8/site-packages/scipy/optimize/_differentiable_functions.py in __init__(self, fun, x0, args, grad, hess, finite_diff_rel_step, finite_diff_bounds, epsilon)
    153 
    154         self._update_grad_impl = update_grad
--> 155         self._update_grad()
    156 
    157         # Hessian Evaluation

~/opt/anaconda3/lib/python3.8/site-packages/scipy/optimize/_differentiable_functions.py in _update_grad(self)
    229     def _update_grad(self):
    230         if not self.g_updated:
--> 231             self._update_grad_impl()
    232             self.g_updated = True
    233 

~/opt/anaconda3/lib/python3.8/site-packages/scipy/optimize/_differentiable_functions.py in update_grad()
    143 
    144             def update_grad():
--> 145                 self.g = grad_wrapped(self.x)
    146 
    147         elif grad in FD_METHODS:

~/opt/anaconda3/lib/python3.8/site-packages/scipy/optimize/_differentiable_functions.py in grad_wrapped(x)
    140             def grad_wrapped(x):
    141                 self.ngev += 1
--> 142                 return np.atleast_1d(grad(x, *args))
    143 
    144             def update_grad():

~/opt/anaconda3/lib/python3.8/site-packages/pennylane/_grad.py in __call__(self, *args, **kwargs)
    115             return ()
    116 
--> 117         grad_value, ans = grad_fn(*args, **kwargs)  # pylint: disable=not-callable
    118         self._forward = ans
    119 

~/opt/anaconda3/lib/python3.8/site-packages/autograd/wrap_util.py in nary_f(*args, **kwargs)
     18             else:
     19                 x = tuple(args[i] for i in argnum)
---> 20             return unary_operator(unary_f, x, *nary_op_args, **nary_op_kwargs)
     21         return nary_f
     22     return nary_operator

~/opt/anaconda3/lib/python3.8/site-packages/pennylane/_grad.py in _grad_with_forward(fun, x)
    133         difference being that it returns both the gradient *and* the forward pass
    134         value."""
--> 135         vjp, ans = _make_vjp(fun, x)
    136 
    137         if not vspace(ans).size == 1:

~/opt/anaconda3/lib/python3.8/site-packages/autograd/core.py in make_vjp(fun, x)
      8 def make_vjp(fun, x):
      9     start_node = VJPNode.new_root()
---> 10     end_value, end_node =  trace(start_node, fun, x)
     11     if end_node is None:
     12         def vjp(g): return vspace(x).zeros()

~/opt/anaconda3/lib/python3.8/site-packages/autograd/tracer.py in trace(start_node, fun, x)
      8     with trace_stack.new_trace() as t:
      9         start_box = new_box(x, t, start_node)
---> 10         end_box = fun(start_box)
     11         if isbox(end_box) and end_box._trace == start_box._trace:
     12             return end_box._value, end_box._node

~/opt/anaconda3/lib/python3.8/site-packages/autograd/wrap_util.py in unary_f(x)
     13                 else:
     14                     subargs = subvals(args, zip(argnum, x))
---> 15                 return fun(*subargs, **kwargs)
     16             if isinstance(argnum, int):
     17                 x = args[argnum]

<ipython-input-68-d9c8029c0ea7> in cost(thetas)
     11     Computes linear entropy from purity.
     12     '''
---> 13     return 1 - purity(thetas)

~/opt/anaconda3/lib/python3.8/site-packages/pennylane/qnode.py in __call__(self, *args, **kwargs)
    851 
    852         # construct the tape
--> 853         self.construct(args, kwargs)
    854 
    855         cache = self.execute_kwargs.get("cache", False)

~/opt/anaconda3/lib/python3.8/site-packages/pennylane/qnode.py in construct(self, args, kwargs)
    755             self.interface = qml.math.get_interface(*args, *list(kwargs.values()))
    756 
--> 757         self._tape = make_qscript(self.func)(*args, **kwargs)
    758         self._qfunc_output = self.tape._qfunc_output
    759 

~/opt/anaconda3/lib/python3.8/site-packages/pennylane/tape/qscript.py in wrapper(*args, **kwargs)
   1376     def wrapper(*args, **kwargs):
   1377         with AnnotatedQueue() as q:
-> 1378             result = fn(*args, **kwargs)
   1379 
   1380         qscript = QuantumScript.from_queue(q)

<ipython-input-68-d9c8029c0ea7> in purity(thetas)
      4     Initializes a density matrix and takes its purity.
      5     '''
----> 6     qml.QubitDensityMatrix(processRho(n, H, tau, thetas, rho), wires=0)
      7     return qml.purity(0)
      8 

~/opt/anaconda3/lib/python3.8/site-packages/pennylane/qnode.py in __call__(self, *args, **kwargs)
    865                 self.execute_kwargs.pop("mode")
    866             # pylint: disable=unexpected-keyword-arg
--> 867             res = qml.execute(
    868                 [self.tape],
    869                 device=self.device,

~/opt/anaconda3/lib/python3.8/site-packages/pennylane/interfaces/execution.py in execute(tapes, device, gradient_fn, interface, grad_on_execution, gradient_kwargs, cache, cachesize, max_diff, override_shots, expand_fn, max_expansion, device_batch_transform)
    405     if gradient_fn == "backprop" or interface is None:
    406         return batch_fn(
--> 407             qml.interfaces.cache_execute(
    408                 batch_execute, cache, return_tuple=False, expand_fn=expand_fn
    409             )(tapes)

~/opt/anaconda3/lib/python3.8/site-packages/pennylane/interfaces/execution.py in wrapper(tapes, **kwargs)
    202         else:
    203             # execute all unique tapes that do not exist in the cache
--> 204             res = fn(execution_tapes.values(), **kwargs)
    205 
    206         final_res = []

~/opt/anaconda3/lib/python3.8/site-packages/pennylane/interfaces/execution.py in fn(tapes, **kwargs)
    128         def fn(tapes: Sequence[QuantumTape], **kwargs):  # pylint: disable=function-redefined
    129             tapes = [expand_fn(tape) for tape in tapes]
--> 130             return original_fn(tapes, **kwargs)
    131 
    132     @wraps(fn)

~/opt/anaconda3/lib/python3.8/contextlib.py in inner(*args, **kwds)
     73         def inner(*args, **kwds):
     74             with self._recreate_cm():
---> 75                 return func(*args, **kwds)
     76         return inner
     77 

~/opt/anaconda3/lib/python3.8/site-packages/pennylane/_qubit_device.py in batch_execute(self, circuits)
    586             self.reset()
    587 
--> 588             res = self.execute(circuit)
    589             results.append(res)
    590 

~/opt/anaconda3/lib/python3.8/site-packages/pennylane/devices/default_mixed.py in execute(self, circuit, **kwargs)
    657                 wires_list.append(m.wires)
    658             self.measured_wires = qml.wires.Wires.all_wires(wires_list)
--> 659         return super().execute(circuit, **kwargs)
    660 
    661     def _execute_legacy(self, circuit, **kwargs):

~/opt/anaconda3/lib/python3.8/site-packages/pennylane/_qubit_device.py in execute(self, circuit, **kwargs)
    316 
    317         # apply all circuit operations
--> 318         self.apply(circuit.operations, rotations=self._get_diagonalizing_gates(circuit), **kwargs)
    319 
    320         # generate computational basis samples

~/opt/anaconda3/lib/python3.8/site-packages/pennylane/devices/default_mixed.py in apply(self, operations, rotations, **kwargs)
    725 
    726         for operation in operations:
--> 727             self._apply_operation(operation)
    728 
    729         # store the pre-rotated state

~/opt/anaconda3/lib/python3.8/site-packages/pennylane/devices/default_mixed.py in _apply_operation(self, operation)
    589             return
    590 
--> 591         matrices = self._get_kraus(operation)
    592 
    593         if operation in diagonal_in_z_basis:

~/opt/anaconda3/lib/python3.8/site-packages/pennylane/devices/default_mixed.py in _get_kraus(self, operation)
    280             return operation.kraus_matrices()
    281 
--> 282         return [operation.matrix()]
    283 
    284     def _apply_channel(self, kraus, wires):

~/opt/anaconda3/lib/python3.8/site-packages/pennylane/operation.py in matrix(self, wire_order)
    743             tensor_like: matrix representation
    744         """
--> 745         canonical_matrix = self.compute_matrix(*self.parameters, **self.hyperparameters)
    746 
    747         if wire_order is None or self.wires == Wires(wire_order):

~/opt/anaconda3/lib/python3.8/site-packages/pennylane/ops/qubit/special_unitary.py in compute_matrix(theta, num_wires)
    447             # jax.numpy.expm does not support broadcasting
    448             return qml.math.stack([qml.math.expm(1j * _A) for _A in A])
--> 449         return qml.math.expm(1j * A)
    450 
    451     def get_one_parameter_generators(self, interface=None):

~/opt/anaconda3/lib/python3.8/site-packages/pennylane/math/multi_dispatch.py in wrapper(*args, **kwargs)
    149             kwargs["like"] = interface
    150 
--> 151             return fn(*args, **kwargs)
    152 
    153         return wrapper

~/opt/anaconda3/lib/python3.8/site-packages/pennylane/math/multi_dispatch.py in expm(tensor, like)
    815     from scipy.linalg import expm as scipy_expm
    816 
--> 817     return scipy_expm(tensor)
    818 
    819 

~/opt/anaconda3/lib/python3.8/site-packages/scipy/linalg/matfuncs.py in expm(A)
    253     # Input checking and conversion is provided by sparse.linalg.expm().
    254     import scipy.sparse.linalg
--> 255     return scipy.sparse.linalg.expm(A)
    256 
    257 

~/opt/anaconda3/lib/python3.8/site-packages/scipy/sparse/linalg/matfuncs.py in expm(A)
    589             [  0.        ,   0.        ,  20.08553692]])
    590     """
--> 591     return _expm(A, use_exact_onenorm='auto')
    592 
    593 

~/opt/anaconda3/lib/python3.8/site-packages/scipy/sparse/linalg/matfuncs.py in _expm(A, use_exact_onenorm)
    638 
    639     # Try Pade order 3.
--> 640     eta_1 = max(h.d4_loose, h.d6_loose)
    641     if eta_1 < 1.495585217958292e-002 and _ell(h.A, 3) == 0:
    642         U, V = h.pade3()

~/opt/anaconda3/lib/python3.8/site-packages/scipy/sparse/linalg/matfuncs.py in d4_loose(self)
    441     def d4_loose(self):
    442         if self.use_exact_onenorm:
--> 443             return self.d4_tight
    444         if self._d4_exact is not None:
    445             return self._d4_exact

~/opt/anaconda3/lib/python3.8/site-packages/scipy/sparse/linalg/matfuncs.py in d4_tight(self)
    417     def d4_tight(self):
    418         if self._d4_exact is None:
--> 419             self._d4_exact = _onenorm(self.A4)**(1/4.)
    420         return self._d4_exact
    421 

~/opt/anaconda3/lib/python3.8/site-packages/scipy/sparse/linalg/matfuncs.py in A4(self)
    390         if self._A4 is None:
    391             self._A4 = _smart_matrix_product(
--> 392                     self.A2, self.A2, structure=self.structure)
    393         return self._A4
    394 

~/opt/anaconda3/lib/python3.8/site-packages/scipy/sparse/linalg/matfuncs.py in A2(self)
    382     def A2(self):
    383         if self._A2 is None:
--> 384             self._A2 = _smart_matrix_product(
    385                     self.A, self.A, structure=self.structure)
    386         return self._A2

~/opt/anaconda3/lib/python3.8/site-packages/scipy/sparse/linalg/matfuncs.py in _smart_matrix_product(A, B, alpha, structure)
    164     else:
    165         if alpha is None:
--> 166             out = A.dot(B)
    167         else:
    168             out = alpha * A.dot(B)

AttributeError: 'ArrayBox' object has no attribute 'dot'

Here is my output from qml.about():

Name: PennyLane
Version: 0.30.0
Summary: PennyLane is a Python quantum machine learning library by Xanadu Inc.
Home-page: https://github.com/XanaduAI/pennylane
Author: 
Author-email: 
License: Apache License 2.0
Location: /Users/carl/opt/anaconda3/lib/python3.8/site-packages
Requires: appdirs, autograd, autoray, cachetools, networkx, numpy, pennylane-lightning, requests, rustworkx, scipy, semantic-version, toml
Required-by: PennyLane-Lightning

Platform info:           macOS-10.16-x86_64-i386-64bit
Python version:          3.8.8
Numpy version:           1.22.0
Scipy version:           1.6.2
Installed devices:
- default.gaussian (PennyLane-0.30.0)
- default.mixed (PennyLane-0.30.0)
- default.qubit (PennyLane-0.30.0)
- default.qubit.autograd (PennyLane-0.30.0)
- default.qubit.jax (PennyLane-0.30.0)
- default.qubit.tf (PennyLane-0.30.0)
- default.qubit.torch (PennyLane-0.30.0)
- default.qutrit (PennyLane-0.30.0)
- null.qubit (PennyLane-0.30.0)
- lightning.qubit (PennyLane-Lightning-0.30.0)

[Solomon]

See these:

AttributeError: 'ArrayBox' object has no attribute 'exp' (very simple code example)

Autograd ArrayBox type intermediate output from optimizer

Thanks @leongfy, I see the issue now This arises because cost_n is computed via a side effect of your cost function. That is, your cost function is updating a global variable, rather than returning the cost. When using Autograd with PennyLane, cost functions must be pure — they cannot perform side effects (such as updating external variables), otherwise: You will see ArrayBox objects rather than NumPy arrays, and The values stored via the side-effect will no longer be different…

[sillygoose]

Sorry, I am not seeing how the linked posts help here. In both the linked cases, the author of the code has explicitly defined a quantity or done something within the definition of their cost function which is not compatible with pennylane such that the ArrayBox error is raised.

In my case, I have only used native pennylane functions and have not modified global variables in the qnodes leading up to my cost function. I did try to use qml.math.toarray() on the final output of my cost function. I also made sure to specify requires_grad=False to each input to the cost function.

[Solomon]

Your stack trace indicates that the exception is thrown from scipy, which is probably being invoked in the minimisation process.

python3.8/site-packages/scipy/sparse/linalg/matfuncs.py in _smart_matrix_product(A, B, alpha, structure)
    164     else:
    165         if alpha is None:
--> 166             out = A.dot(B)
    167         else:
    168             out = alpha * A.dot(B)

AttributeError: 'ArrayBox' object has no attribute 'dot'

I suggested the first link since they are using autograd which you may wish to explore as an alternative, and the second uses the qml build in optimiser.
BTW, I am trying to help you as a fellow programmer, I am not from the Xanadu support team.

[isaacdevlugt]

Hey @sillygoose! Welcome to the forum

I managed to narrow in on what’s going on and it seems like your error stems from something akin to this example:

dev = qml.device("default.qubit", wires=2)


@qml.qnode(dev)
def circuit_sun(x):
    qml.SpecialUnitary(x, wires=dev.wires)
    qml.CNOT(wires=[0, 1])
    return qml.expval(qml.PauliZ(0))


@qml.qnode(dev)
def circuit_rx(x):
    qml.RX(x[0], wires=0)
    qml.RX(x[1], wires=1)
    qml.CNOT(wires=[0, 1])
    return qml.expval(qml.PauliZ(0))


def cost_sun(x):
    return np.abs(circuit_sun(x) - 0.5) ** 2

def cost_rx(x):
    return np.abs(circuit_rx(x) - 0.5) ** 2

x_sun = np.random.uniform(0, 1, size=(15,), requires_grad=True)
x_rx = np.random.uniform(0, 1, size=(2,), requires_grad=True)

minimize(cost_rx, x_rx, method="BFGS", jac=qml.grad(cost_rx, argnum=0)) # works
minimize(cost_sun, x_sun, method="BFGS", jac=qml.grad(cost_sun, argnum=0)) # doesn't work

Seems to be an issue with SpecialUnitary. Curiously, without the jac keyword argument everything works fine. I’m going to check in with someone on the development team and I will get back to you ASAP!

[sillygoose]

Thank you very much :)!

[Wierichs]

Hi @sillygoose,

The problem indeed is with SpecialUnitary. Because it uses the function scipy.linalg.expm, which is not differentiable in Autograd, differentiation of SpecialUnitary with automatic differentiation (as is used here by default in your and @isaacdevlugt’s examples) is not supported. Also see the docs for more information on the operation.
The reason for the MWE running if you skip the jac keyword is that - as far as I know - scipy will default to using finite differences to compute the gradient if you don’t pass a method to do so. I’d not recommend this, in particular for larger examples.
Alternatively, I’d recommend to

1. Move to one of the other machine learning frameworks: JAX, Tensorflow, Pytorch. They all support differentiation of SpecialUnitary (via support for expm).
2. Continue to pass the jac keyword argument explicitly, to avoid finite differences.

The following should work

import pennylane as qml
import numpy as np
from jax import numpy as jnp
import jax
jax.config.update("jax_enable_x64", True)
from scipy.optimize import minimize

n = 3 # total number of qubits
d = 2**n # dimension of composite system

coeffs = [0.5]*3
ops = [qml.PauliZ(0) @ qml.PauliZ(1), qml.PauliZ(1) @ qml.PauliZ(2), qml.PauliZ(0) @ qml.PauliZ(2)]
isingHam = qml.Hamiltonian(coeffs, ops)
H = isingHam # native Hamiltonian
tau = 1 # characteristic time

rho = np.zeros((d, d), dtype=np.complex128)
val = 0.25
rho[0, 0] = val
rho[1, 1] = val
rho[2, 2] = val
rho[3, 3] = val
rho = np.array(rho) # initial density matrix

wireList = list(range(n))
thetas = jnp.array(np.random.randn((d**2-1)))

devRho = qml.device("default.mixed", wires=3)
@qml.qnode(devRho)
def processRho(n, H, tau, thetas, rho):
    '''
    Time evolves rho using scrambled Hamiltonian by characteristic time, then traces out of environment.
    Inputs:
     - n: total number of qubits
     - H: native Hamiltonian to be scrambled
     - tau: characteristic time
     - thetas: parameters to scramble H
     - rho: initial density matrix
    
    Outputs:
     - rhoTau: time evolved, then traced out density matrix
    '''
    qml.QubitDensityMatrix(rho, wires=wireList)
    qml.adjoint(qml.SpecialUnitary(thetas, wires=wireList))
    qml.ApproxTimeEvolution(H, tau, 100)
    qml.SpecialUnitary(thetas, wires=wireList)
    return qml.density_matrix([0])

@qml.qnode(devRho)
def purity(thetas):
    '''
    Initializes a density matrix and takes its purity.
    '''
    qml.QubitDensityMatrix(processRho(n, H, tau, thetas, rho), wires=0)
    return qml.purity(0)

def cost(thetas):
    '''
    Computes linear entropy from purity. 
    '''
    return 1 - purity(thetas)

thetas = jnp.array(thetas)

minimize(purity, thetas, method='BFGS', jac=jax.grad(cost, argnums=0))

Hope this helps! Happy coding