Postprocess jacobian result shape

[albi3ro]

The gradients also need a different result under the new return types designation. This change handles that post-processing.

According to @lillian542 , there are still some test failures due to use of qml.jacobian and an extra dimension occurring somewhere that is not getting squeezed out, but this change should fix at least one test failure.

[codecov]

Codecov Report

Merging #159 (4b42335) into main (124dd9e) will not change coverage.
The diff coverage is 100.00%.

@@            Coverage Diff            @@
##              main      #159   +/-   ##
=========================================
  Coverage   100.00%   100.00%           
=========================================
  Files            7         7           
  Lines          957       957           
  Branches       228       228           
=========================================
  Hits           957       957           

Impacted Files	Coverage Δ
src/braket/pennylane_plugin/braket_device.py	100.00% <100.00%> (ø)

[speller26]

Integ test failures

=================================== FAILURES ===================================
_______________ TestExpval.test_paulix_expectation[device0-None] _______________
[gw0] linux -- Python 3.8.16 /root/.pyenv/versions/3.8.16/bin/python3.8

self = <test_expval.TestExpval object at 0x7f9b1f75f340>
device = <function device.<locals>._device at 0x7f9b1a333a60>, shots = None
tol = {'atol': 0.05, 'rtol': 0.1}

    def test_paulix_expectation(self, device, shots, tol):
        """Test that PauliX expectation value is correct"""
        dev = device(2)
    
        theta = 0.432
        phi = 0.123
    
        @qml.qnode(dev)
        def circuit():
            qml.RY(theta, wires=[0])
            qml.RY(phi, wires=[1])
            qml.CNOT(wires=[0, 1])
            return [qml.expval(qml.PauliX(i)) for i in range(2)]
    
>       assert np.allclose(circuit(), np.array([np.sin(theta) * np.sin(phi), np.sin(phi)]), **tol)

test/integ_tests/test_expval.py:74: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 

self = <QNode: wires=2, device='braket.aws.qubit', interface='auto', diff_method='best'>
args = (), kwargs = {}, override_shots = False, old_interface = 'auto'
cache = True, using_custom_cache = False, res = array(0.05136886)

    def __call__(self, *args, **kwargs):  # pylint: disable=too-many-branches, too-many-statements
        override_shots = False
        old_interface = self.interface
    
        if old_interface == "auto":
            self.interface = qml.math.get_interface(*args, *list(kwargs.values()))
            self.device.tracker = self._original_device.tracker
    
        if not self._qfunc_uses_shots_arg:
            # If shots specified in call but not in qfunc signature,
            # interpret it as device shots value for this call.
            override_shots = kwargs.pop("shots", False)
    
            if override_shots is not False:
                # Since shots has changed, we need to update the preferred gradient function.
                # This is because the gradient function chosen at initialization may
                # no longer be applicable.
    
                # store the initialization gradient function
                original_grad_fn = [self.gradient_fn, self.gradient_kwargs, self.device]
    
                # pylint: disable=not-callable
                # update the gradient function
                set_shots(self._original_device, override_shots)(self._update_gradient_fn)()
    
        # construct the tape
        self.construct(args, kwargs)
    
        cache = self.execute_kwargs.get("cache", False)
        using_custom_cache = (
            hasattr(cache, "__getitem__")
            and hasattr(cache, "__setitem__")
            and hasattr(cache, "__delitem__")
        )
        self._tape_cached = using_custom_cache and self.tape.hash in cache
    
        if qml.active_return():
            if "mode" in self.execute_kwargs:
                self.execute_kwargs.pop("mode")
            # pylint: disable=unexpected-keyword-arg
            res = qml.execute(
                [self.tape],
                device=self.device,
                gradient_fn=self.gradient_fn,
                interface=self.interface,
                gradient_kwargs=self.gradient_kwargs,
                override_shots=override_shots,
                **self.execute_kwargs,
            )
    
            res = res[0]
    
            if old_interface == "auto":
                self.interface = "auto"
    
            # Special case of single Measurement in a list
            if isinstance(self._qfunc_output, list) and len(self._qfunc_output) == 1:
                return [res]
    
            # If the return type is not tuple (list or ndarray) (Autograd and TF backprop removed)
            if not isinstance(self._qfunc_output, (tuple, qml.measurements.MeasurementProcess)):
                if self.device._shot_vector:
                    res = [type(self.tape._qfunc_output)(r) for r in res]
                    res = tuple(res)
    
                else:
>                   res = type(self.tape._qfunc_output)(res)
E                   TypeError: iteration over a 0-d array

/root/.pyenv/versions/3.8.16/lib/python3.8/site-packages/pennylane/qnode.py:893: TypeError
_______________ TestExpval.test_pauliy_expectation[device0-None] _______________
[gw0] linux -- Python 3.8.16 /root/.pyenv/versions/3.8.16/bin/python3.8

self = <test_expval.TestExpval object at 0x7f9b2e0067c0>
device = <function device.<locals>._device at 0x7f9b17a9b820>, shots = None
tol = {'atol': 0.05, 'rtol': 0.1}

    def test_pauliy_expectation(self, device, shots, tol):
        """Test that PauliY expectation value is correct"""
        dev = device(2)
    
        theta = 0.432
        phi = 0.123
    
        @qml.qnode(dev)
        def circuit():
            qml.RX(theta, wires=[0])
            qml.RX(phi, wires=[1])
            qml.CNOT(wires=[0, 1])
            return [qml.expval(qml.PauliY(i)) for i in range(2)]
    
>       assert np.allclose(circuit(), np.array([0, -np.cos(theta) * np.sin(phi)]), **tol)

test/integ_tests/test_expval.py:90: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 

self = <QNode: wires=2, device='braket.aws.qubit', interface='auto', diff_method='best'>
args = (), kwargs = {}, override_shots = False, old_interface = 'auto'
cache = True, using_custom_cache = False, res = array(0.)

    def __call__(self, *args, **kwargs):  # pylint: disable=too-many-branches, too-many-statements
        override_shots = False
        old_interface = self.interface
    
        if old_interface == "auto":
            self.interface = qml.math.get_interface(*args, *list(kwargs.values()))
            self.device.tracker = self._original_device.tracker
    
        if not self._qfunc_uses_shots_arg:
            # If shots specified in call but not in qfunc signature,
            # interpret it as device shots value for this call.
            override_shots = kwargs.pop("shots", False)
    
            if override_shots is not False:
                # Since shots has changed, we need to update the preferred gradient function.
                # This is because the gradient function chosen at initialization may
                # no longer be applicable.
    
                # store the initialization gradient function
                original_grad_fn = [self.gradient_fn, self.gradient_kwargs, self.device]
    
                # pylint: disable=not-callable
                # update the gradient function
                set_shots(self._original_device, override_shots)(self._update_gradient_fn)()
    
        # construct the tape
        self.construct(args, kwargs)
    
        cache = self.execute_kwargs.get("cache", False)
        using_custom_cache = (
            hasattr(cache, "__getitem__")
            and hasattr(cache, "__setitem__")
            and hasattr(cache, "__delitem__")
        )
        self._tape_cached = using_custom_cache and self.tape.hash in cache
    
        if qml.active_return():
            if "mode" in self.execute_kwargs:
                self.execute_kwargs.pop("mode")
            # pylint: disable=unexpected-keyword-arg
            res = qml.execute(
                [self.tape],
                device=self.device,
                gradient_fn=self.gradient_fn,
                interface=self.interface,
                gradient_kwargs=self.gradient_kwargs,
                override_shots=override_shots,
                **self.execute_kwargs,
            )
    
            res = res[0]
    
            if old_interface == "auto":
                self.interface = "auto"
    
            # Special case of single Measurement in a list
            if isinstance(self._qfunc_output, list) and len(self._qfunc_output) == 1:
                return [res]
    
            # If the return type is not tuple (list or ndarray) (Autograd and TF backprop removed)
            if not isinstance(self._qfunc_output, (tuple, qml.measurements.MeasurementProcess)):
                if self.device._shot_vector:
                    res = [type(self.tape._qfunc_output)(r) for r in res]
                    res = tuple(res)
    
                else:
>                   res = type(self.tape._qfunc_output)(res)
E                   TypeError: iteration over a 0-d array

/root/.pyenv/versions/3.8.16/lib/python3.8/site-packages/pennylane/qnode.py:893: TypeError
______________ TestExpval.test_hadamard_expectation[device0-None] ______________
[gw0] linux -- Python 3.8.16 /root/.pyenv/versions/3.8.16/bin/python3.8

self = <test_expval.TestExpval object at 0x7f9b1f75f1c0>
device = <function device.<locals>._device at 0x7f9b1a111670>, shots = None
tol = {'atol': 0.05, 'rtol': 0.1}

    def test_hadamard_expectation(self, device, shots, tol):
        """Test that Hadamard expectation value is correct"""
        dev = device(2)
    
        theta = 0.432
        phi = 0.123
    
        @qml.qnode(dev)
        def circuit():
            qml.RY(theta, wires=[0])
            qml.RY(phi, wires=[1])
            qml.CNOT(wires=[0, 1])
            return [qml.expval(qml.Hadamard(i)) for i in range(2)]
    
        expected = np.array(
            [np.sin(theta) * np.sin(phi) + np.cos(theta), np.cos(theta) * np.cos(phi) + np.sin(phi)]
        ) / np.sqrt(2)
>       assert np.allclose(circuit(), expected, **tol)

test/integ_tests/test_expval.py:109: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 

self = <QNode: wires=2, device='braket.aws.qubit', interface='auto', diff_method='best'>
args = (), kwargs = {}, override_shots = False, old_interface = 'auto'
cache = True, using_custom_cache = False, res = array(0.67846828)

    def __call__(self, *args, **kwargs):  # pylint: disable=too-many-branches, too-many-statements
        override_shots = False
        old_interface = self.interface
    
        if old_interface == "auto":
            self.interface = qml.math.get_interface(*args, *list(kwargs.values()))
            self.device.tracker = self._original_device.tracker
    
        if not self._qfunc_uses_shots_arg:
            # If shots specified in call but not in qfunc signature,
            # interpret it as device shots value for this call.
            override_shots = kwargs.pop("shots", False)
    
            if override_shots is not False:
                # Since shots has changed, we need to update the preferred gradient function.
                # This is because the gradient function chosen at initialization may
                # no longer be applicable.
    
                # store the initialization gradient function
                original_grad_fn = [self.gradient_fn, self.gradient_kwargs, self.device]
    
                # pylint: disable=not-callable
                # update the gradient function
                set_shots(self._original_device, override_shots)(self._update_gradient_fn)()
    
        # construct the tape
        self.construct(args, kwargs)
    
        cache = self.execute_kwargs.get("cache", False)
        using_custom_cache = (
            hasattr(cache, "__getitem__")
            and hasattr(cache, "__setitem__")
            and hasattr(cache, "__delitem__")
        )
        self._tape_cached = using_custom_cache and self.tape.hash in cache
    
        if qml.active_return():
            if "mode" in self.execute_kwargs:
                self.execute_kwargs.pop("mode")
            # pylint: disable=unexpected-keyword-arg
            res = qml.execute(
                [self.tape],
                device=self.device,
                gradient_fn=self.gradient_fn,
                interface=self.interface,
                gradient_kwargs=self.gradient_kwargs,
                override_shots=override_shots,
                **self.execute_kwargs,
            )
    
            res = res[0]
    
            if old_interface == "auto":
                self.interface = "auto"
    
            # Special case of single Measurement in a list
            if isinstance(self._qfunc_output, list) and len(self._qfunc_output) == 1:
                return [res]
    
            # If the return type is not tuple (list or ndarray) (Autograd and TF backprop removed)
            if not isinstance(self._qfunc_output, (tuple, qml.measurements.MeasurementProcess)):
                if self.device._shot_vector:
                    res = [type(self.tape._qfunc_output)(r) for r in res]
                    res = tuple(res)
    
                else:
>                   res = type(self.tape._qfunc_output)(res)
E                   TypeError: iteration over a 0-d array

/root/.pyenv/versions/3.8.16/lib/python3.8/site-packages/pennylane/qnode.py:893: TypeError
_____________ TestExpval.test_hermitian_expectation[device0-None] ______________
[gw0] linux -- Python 3.8.16 /root/.pyenv/versions/3.8.16/bin/python3.8

self = <test_expval.TestExpval object at 0x7f9b2e006370>
device = <function device.<locals>._device at 0x7f9b19565160>, shots = None
tol = {'atol': 0.05, 'rtol': 0.1}

    def test_hermitian_expectation(self, device, shots, tol):
        """Test that arbitrary Hermitian expectation values are correct"""
        dev = device(2)
    
        theta = 0.432
        phi = 0.123
    
        @qml.qnode(dev)
        def circuit():
            qml.RY(theta, wires=[0])
            qml.RY(phi, wires=[1])
            qml.CNOT(wires=[0, 1])
            return [qml.expval(qml.Hermitian(A, i)) for i in range(2)]
    
        a = A[0, 0]
        re_b = A[0, 1].real
        d = A[1, 1]
        ev1 = ((a - d) * np.cos(theta) + 2 * re_b * np.sin(theta) * np.sin(phi) + a + d) / 2
        ev2 = ((a - d) * np.cos(theta) * np.cos(phi) + 2 * re_b * np.sin(phi) + a + d) / 2
        expected = np.array([ev1, ev2])
    
>       assert np.allclose(circuit(), expected, **tol)

test/integ_tests/test_expval.py:132: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 

self = <QNode: wires=2, device='braket.aws.qubit', interface='auto', diff_method='best'>
args = (), kwargs = {}, override_shots = False, old_interface = 'auto'
cache = True, using_custom_cache = False, res = array(1.12045643)

    def __call__(self, *args, **kwargs):  # pylint: disable=too-many-branches, too-many-statements
        override_shots = False
        old_interface = self.interface
    
        if old_interface == "auto":
            self.interface = qml.math.get_interface(*args, *list(kwargs.values()))
            self.device.tracker = self._original_device.tracker
    
        if not self._qfunc_uses_shots_arg:
            # If shots specified in call but not in qfunc signature,
            # interpret it as device shots value for this call.
            override_shots = kwargs.pop("shots", False)
    
            if override_shots is not False:
                # Since shots has changed, we need to update the preferred gradient function.
                # This is because the gradient function chosen at initialization may
                # no longer be applicable.
    
                # store the initialization gradient function
                original_grad_fn = [self.gradient_fn, self.gradient_kwargs, self.device]
    
                # pylint: disable=not-callable
                # update the gradient function
                set_shots(self._original_device, override_shots)(self._update_gradient_fn)()
    
        # construct the tape
        self.construct(args, kwargs)
    
        cache = self.execute_kwargs.get("cache", False)
        using_custom_cache = (
            hasattr(cache, "__getitem__")
            and hasattr(cache, "__setitem__")
            and hasattr(cache, "__delitem__")
        )
        self._tape_cached = using_custom_cache and self.tape.hash in cache
    
        if qml.active_return():
            if "mode" in self.execute_kwargs:
                self.execute_kwargs.pop("mode")
            # pylint: disable=unexpected-keyword-arg
            res = qml.execute(
                [self.tape],
                device=self.device,
                gradient_fn=self.gradient_fn,
                interface=self.interface,
                gradient_kwargs=self.gradient_kwargs,
                override_shots=override_shots,
                **self.execute_kwargs,
            )
    
            res = res[0]
    
            if old_interface == "auto":
                self.interface = "auto"
    
            # Special case of single Measurement in a list
            if isinstance(self._qfunc_output, list) and len(self._qfunc_output) == 1:
                return [res]
    
            # If the return type is not tuple (list or ndarray) (Autograd and TF backprop removed)
            if not isinstance(self._qfunc_output, (tuple, qml.measurements.MeasurementProcess)):
                if self.device._shot_vector:
                    res = [type(self.tape._qfunc_output)(r) for r in res]
                    res = tuple(res)
    
                else:
>                   res = type(self.tape._qfunc_output)(res)
E                   TypeError: iteration over a 0-d array

/root/.pyenv/versions/3.8.16/lib/python3.8/site-packages/pennylane/qnode.py:893: TypeError
__________________ TestExpval.test_hamiltonian[device0-None] ___________________
[gw0] linux -- Python 3.8.16 /root/.pyenv/versions/3.8.16/bin/python3.8

self = <test_expval.TestExpval object at 0x7f9b1f706d30>
device = <function device.<locals>._device at 0x7f9b1a111430>, shots = None
tol = {'atol': 0.05, 'rtol': 0.1}

    def test_hamiltonian(self, device, shots, tol):
        """Test that Hamiltonian expectation value is correct"""
        dev = device(2)
    
        theta = 0.432
        phi = 0.123
    
        @qml.qnode(dev)
        def circuit():
            qml.RX(theta, wires=[0])
            qml.RX(phi, wires=[1])
            qml.CNOT(wires=[0, 1])
            return qml.expval(qml.Hamiltonian((2, 3), (qml.PauliZ(0), qml.PauliY(1))))
    
>       assert np.allclose(circuit(), 2 * np.cos(theta) - 3 * np.cos(theta) * np.sin(phi), **tol)

test/integ_tests/test_expval.py:181: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 
/root/.pyenv/versions/3.8.16/lib/python3.8/site-packages/pennylane/qnode.py:867: in __call__
    res = qml.execute(
/root/.pyenv/versions/3.8.16/lib/python3.8/site-packages/pennylane/interfaces/execution.py:493: in execute
    return batch_fn(res)
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 

res = array(1.48200466)

    def processing_fn(res):
        count = 0
        final_results = []
    
        for idx, s in enumerate(tape_counts):
            # apply any batch transform post-processing
>           new_res = batch_fns[idx](res[count : count + s])
E           IndexError: too many indices for array: array is 0-dimensional, but 1 were indexed

/root/.pyenv/versions/3.8.16/lib/python3.8/site-packages/pennylane/transforms/batch_transform.py:510: IndexError
_________________ TestExpval.test_nondiff_param[device0-None] __________________
[gw0] linux -- Python 3.8.16 /root/.pyenv/versions/3.8.16/bin/python3.8

self = <test_expval.TestExpval object at 0x7f9b1f75fa30>
device = <function device.<locals>._device at 0x7f9b1ec509d0>

    def test_nondiff_param(self, device):
        """Test that the device can be successfully executed with the Autograd
        interface when some of the arguments are tensors with requires_grad=False"""
    
        dev1 = device(1)
        dev2 = qml.device("default.qubit", wires=1)
    
        if not dev1.analytic:
            pytest.skip("This test is designed to work in analytic mode")
    
        def circuit(x, y):
            qml.RX(x[0], wires=0)
            qml.Rot(*x[1:], wires=0)
            qml.RY(y[0], wires=0)
            return qml.expval(qml.PauliZ(0))
    
        x = np.array([0.1, 0.2, 0.3, 0.4], requires_grad=False)
        y = np.array([0.5], requires_grad=True)
    
        qnode1 = qml.QNode(circuit, dev1, interface="autograd")
        qnode2 = qml.QNode(circuit, dev2, interface="autograd")
    
>       r1 = qnode1(x, y)

test/integ_tests/test_expval.py:205: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 
/root/.pyenv/versions/3.8.16/lib/python3.8/site-packages/pennylane/qnode.py:867: in __call__
    res = qml.execute(
/root/.pyenv/versions/3.8.16/lib/python3.8/site-packages/pennylane/interfaces/execution.py:493: in execute
    return batch_fn(res)
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 

res = 0.6725763684370558

    def processing_fn(res):
        count = 0
        final_results = []
    
        for idx, s in enumerate(tape_counts):
            # apply any batch transform post-processing
>           new_res = batch_fns[idx](res[count : count + s])
E           IndexError: invalid index to scalar variable.

/root/.pyenv/versions/3.8.16/lib/python3.8/site-packages/pennylane/transforms/batch_transform.py:510: IndexError
______________ TestTensorExpval.test_paulix_pauliy[device0-None] _______________
[gw0] linux -- Python 3.8.16 /root/.pyenv/versions/3.8.16/bin/python3.8

self = <test_expval.TestTensorExpval object at 0x7f9b1f708760>
device = <function device.<locals>._device at 0x7f9b1ec6c280>, shots = None
tol = {'atol': 0.05, 'rtol': 0.1}

    def test_paulix_pauliy(self, device, shots, tol):
        """Test that a tensor product involving PauliX and PauliY works correctly"""
        dev = device(3)
    
        theta = 0.432
        phi = 0.123
        varphi = -0.543
    
        @qml.qnode(dev)
        def circuit():
            qml.RX(theta, wires=[0])
            qml.RX(phi, wires=[1])
            qml.RX(varphi, wires=[2])
            qml.CNOT(wires=[0, 1])
            qml.CNOT(wires=[1, 2])
            return qml.expval(qml.PauliX(0) @ qml.PauliY(2))
    
>       assert np.allclose(circuit(), np.sin(theta) * np.sin(phi) * np.sin(varphi), **tol)

test/integ_tests/test_expval.py:232: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 
/root/.pyenv/versions/3.8.16/lib/python3.8/site-packages/pennylane/qnode.py:867: in __call__
    res = qml.execute(
/root/.pyenv/versions/3.8.16/lib/python3.8/site-packages/pennylane/interfaces/execution.py:493: in execute
    return batch_fn(res)
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 

res = array(-0.02654264)

    def processing_fn(res):
        count = 0
        final_results = []
    
        for idx, s in enumerate(tape_counts):
            # apply any batch transform post-processing
>           new_res = batch_fns[idx](res[count : count + s])
E           IndexError: too many indices for array: array is 0-dimensional, but 1 were indexed

/root/.pyenv/versions/3.8.16/lib/python3.8/site-packages/pennylane/transforms/batch_transform.py:510: IndexError
_____________ TestTensorExpval.test_pauliz_hadamard[device0-None] ______________
[gw0] linux -- Python 3.8.16 /root/.pyenv/versions/3.8.16/bin/python3.8

self = <test_expval.TestTensorExpval object at 0x7f9b1f708e80>
device = <function device.<locals>._device at 0x7f9b120684c0>, shots = None
tol = {'atol': 0.05, 'rtol': 0.1}

    def test_pauliz_hadamard(self, device, shots, tol):
        """Test that a tensor product involving PauliZ and PauliY and Hadamard works correctly"""
        dev = device(3)
    
        theta = 0.432
        phi = 0.123
        varphi = -0.543
    
        @qml.qnode(dev)
        def circuit():
            qml.RX(theta, wires=[0])
            qml.RX(phi, wires=[1])
            qml.RX(varphi, wires=[2])
            qml.CNOT(wires=[0, 1])
            qml.CNOT(wires=[1, 2])
            return qml.expval(
                qml.PauliZ(wires=[0]) @ qml.Hadamard(wires=[1]) @ qml.PauliY(wires=[2])
            )
    
        expected = -(np.cos(varphi) * np.sin(phi) + np.sin(varphi) * np.cos(theta)) / np.sqrt(2)
>       assert np.allclose(circuit(), expected, **tol)

test/integ_tests/test_expval.py:254: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 
/root/.pyenv/versions/3.8.16/lib/python3.8/site-packages/pennylane/qnode.py:867: in __call__
    res = qml.execute(
/root/.pyenv/versions/3.8.16/lib/python3.8/site-packages/pennylane/interfaces/execution.py:493: in execute
    return batch_fn(res)
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 

res = array(0.25752433)

    def processing_fn(res):
        count = 0
        final_results = []
    
        for idx, s in enumerate(tape_counts):
            # apply any batch transform post-processing
>           new_res = batch_fns[idx](res[count : count + s])
E           IndexError: too many indices for array: array is 0-dimensional, but 1 were indexed

/root/.pyenv/versions/3.8.16/lib/python3.8/site-packages/pennylane/transforms/batch_transform.py:510: IndexError
______________ TestExpval.test_identity_expectation[device0-None] ______________
[gw1] linux -- Python 3.8.16 /root/.pyenv/versions/3.8.16/bin/python3.8

self = <test_expval.TestExpval object at 0x7fcd9ee9b3a0>
device = <function device.<locals>._device at 0x7fcd8500a160>, shots = None
tol = {'atol': 0.05, 'rtol': 0.1}

    def test_identity_expectation(self, device, shots, tol):
        """Test that identity expectation value (i.e. the trace) is 1"""
        dev = device(2)
    
        theta = 0.432
        phi = 0.123
    
        @qml.qnode(dev)
        def circuit():
            qml.RX(theta, wires=[0])
            qml.RX(phi, wires=[1])
            qml.CNOT(wires=[0, 1])
            return [qml.expval(qml.Identity(i)) for i in range(2)]
    
>       assert np.allclose(circuit(), np.array([1, 1]), **tol)

test/integ_tests/test_expval.py:42: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 

self = <QNode: wires=2, device='braket.aws.qubit', interface='auto', diff_method='best'>
args = (), kwargs = {}, override_shots = False, old_interface = 'auto'
cache = True, using_custom_cache = False, res = array(1.)

    def __call__(self, *args, **kwargs):  # pylint: disable=too-many-branches, too-many-statements
        override_shots = False
        old_interface = self.interface
    
        if old_interface == "auto":
            self.interface = qml.math.get_interface(*args, *list(kwargs.values()))
            self.device.tracker = self._original_device.tracker
    
        if not self._qfunc_uses_shots_arg:
            # If shots specified in call but not in qfunc signature,
            # interpret it as device shots value for this call.
            override_shots = kwargs.pop("shots", False)
    
            if override_shots is not False:
                # Since shots has changed, we need to update the preferred gradient function.
                # This is because the gradient function chosen at initialization may
                # no longer be applicable.
    
                # store the initialization gradient function
                original_grad_fn = [self.gradient_fn, self.gradient_kwargs, self.device]
    
                # pylint: disable=not-callable
                # update the gradient function
                set_shots(self._original_device, override_shots)(self._update_gradient_fn)()
    
        # construct the tape
        self.construct(args, kwargs)
    
        cache = self.execute_kwargs.get("cache", False)
        using_custom_cache = (
            hasattr(cache, "__getitem__")
            and hasattr(cache, "__setitem__")
            and hasattr(cache, "__delitem__")
        )
        self._tape_cached = using_custom_cache and self.tape.hash in cache
    
        if qml.active_return():
            if "mode" in self.execute_kwargs:
                self.execute_kwargs.pop("mode")
            # pylint: disable=unexpected-keyword-arg
            res = qml.execute(
                [self.tape],
                device=self.device,
                gradient_fn=self.gradient_fn,
                interface=self.interface,
                gradient_kwargs=self.gradient_kwargs,
                override_shots=override_shots,
                **self.execute_kwargs,
            )
    
            res = res[0]
    
            if old_interface == "auto":
                self.interface = "auto"
    
            # Special case of single Measurement in a list
            if isinstance(self._qfunc_output, list) and len(self._qfunc_output) == 1:
                return [res]
    
            # If the return type is not tuple (list or ndarray) (Autograd and TF backprop removed)
            if not isinstance(self._qfunc_output, (tuple, qml.measurements.MeasurementProcess)):
                if self.device._shot_vector:
                    res = [type(self.tape._qfunc_output)(r) for r in res]
                    res = tuple(res)
    
                else:
>                   res = type(self.tape._qfunc_output)(res)
E                   TypeError: iteration over a 0-d array

/root/.pyenv/versions/3.8.16/lib/python3.8/site-packages/pennylane/qnode.py:893: TypeError
_______________ TestExpval.test_pauliz_expectation[device0-None] _______________
[gw1] linux -- Python 3.8.16 /root/.pyenv/versions/3.8.16/bin/python3.8

self = <test_expval.TestExpval object at 0x7fcd9ee9bdc0>
device = <function device.<locals>._device at 0x7fcd8500aee0>, shots = None
tol = {'atol': 0.05, 'rtol': 0.1}

    def test_pauliz_expectation(self, device, shots, tol):
        """Test that PauliZ expectation value is correct"""
        dev = device(2)
    
        theta = 0.432
        phi = 0.123
    
        @qml.qnode(dev)
        def circuit():
            qml.RX(theta, wires=[0])
            qml.RX(phi, wires=[1])
            qml.CNOT(wires=[0, 1])
            return [qml.expval(qml.PauliZ(i)) for i in range(2)]
    
>       assert np.allclose(circuit(), np.array([np.cos(theta), np.cos(theta) * np.cos(phi)]), **tol)

test/integ_tests/test_expval.py:58: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 

self = <QNode: wires=2, device='braket.aws.qubit', interface='auto', diff_method='best'>
args = (), kwargs = {}, override_shots = False, old_interface = 'auto'
cache = True, using_custom_cache = False, res = array(0.90813019)

    def __call__(self, *args, **kwargs):  # pylint: disable=too-many-branches, too-many-statements
        override_shots = False
        old_interface = self.interface
    
        if old_interface == "auto":
            self.interface = qml.math.get_interface(*args, *list(kwargs.values()))
            self.device.tracker = self._original_device.tracker
    
        if not self._qfunc_uses_shots_arg:
            # If shots specified in call but not in qfunc signature,
            # interpret it as device shots value for this call.
            override_shots = kwargs.pop("shots", False)
    
            if override_shots is not False:
                # Since shots has changed, we need to update the preferred gradient function.
                # This is because the gradient function chosen at initialization may
                # no longer be applicable.
    
                # store the initialization gradient function
                original_grad_fn = [self.gradient_fn, self.gradient_kwargs, self.device]
    
                # pylint: disable=not-callable
                # update the gradient function
                set_shots(self._original_device, override_shots)(self._update_gradient_fn)()
    
        # construct the tape
        self.construct(args, kwargs)
    
        cache = self.execute_kwargs.get("cache", False)
        using_custom_cache = (
            hasattr(cache, "__getitem__")
            and hasattr(cache, "__setitem__")
            and hasattr(cache, "__delitem__")
        )
        self._tape_cached = using_custom_cache and self.tape.hash in cache
    
        if qml.active_return():
            if "mode" in self.execute_kwargs:
                self.execute_kwargs.pop("mode")
            # pylint: disable=unexpected-keyword-arg
            res = qml.execute(
                [self.tape],
                device=self.device,
                gradient_fn=self.gradient_fn,
                interface=self.interface,
                gradient_kwargs=self.gradient_kwargs,
                override_shots=override_shots,
                **self.execute_kwargs,
            )
    
            res = res[0]
    
            if old_interface == "auto":
                self.interface = "auto"
    
            # Special case of single Measurement in a list
            if isinstance(self._qfunc_output, list) and len(self._qfunc_output) == 1:
                return [res]
    
            # If the return type is not tuple (list or ndarray) (Autograd and TF backprop removed)
            if not isinstance(self._qfunc_output, (tuple, qml.measurements.MeasurementProcess)):
                if self.device._shot_vector:
                    res = [type(self.tape._qfunc_output)(r) for r in res]
                    res = tuple(res)
    
                else:
>                   res = type(self.tape._qfunc_output)(res)
E                   TypeError: iteration over a 0-d array

/root/.pyenv/versions/3.8.16/lib/python3.8/site-packages/pennylane/qnode.py:893: TypeError
____________ TestDeviceIntegration.test_one_qubit_circuit[None-d0] _____________
[gw0] linux -- Python 3.8.16 /root/.pyenv/versions/3.8.16/bin/python3.8

self = <test_integration.TestDeviceIntegration object at 0x7f9b1f6ee550>
shots = None
d = ('braket.aws.qubit', 'arn:aws:braket:::device/quantum-simulator/amazon/sv1')
tol = {'atol': 0.05, 'rtol': 0.1}
extra_kwargs = <function extra_kwargs.<locals>._extra_kwargs at 0x7f9b1718a160>

    @pytest.mark.parametrize("d", shortname_and_backends)
    @pytest.mark.parametrize("shots", [None, 8192])
    def test_one_qubit_circuit(self, shots, d, tol, extra_kwargs):
        """Test that devices provide correct result for a simple circuit"""
        dev = TestDeviceIntegration._device(d, 1, extra_kwargs)
    
        a = 0.543
        b = 0.123
        c = 0.987
    
        @qml.qnode(dev)
        def circuit(x, y, z):
            """Reference QNode"""
            qml.BasisState(np.array([1]), wires=0)
            qml.Hadamard(wires=0)
            qml.Rot(x, y, z, wires=0)
            return qml.expval(qml.PauliZ(0))
    
>       assert np.allclose(circuit(a, b, c), np.cos(a) * np.sin(b), **tol)

test/integ_tests/test_integration.py:64: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 
/root/.pyenv/versions/3.8.16/lib/python3.8/site-packages/pennylane/qnode.py:867: in __call__
    res = qml.execute(
/root/.pyenv/versions/3.8.16/lib/python3.8/site-packages/pennylane/interfaces/execution.py:493: in execute
    return batch_fn(res)
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 

res = array(0.10504264)

    def processing_fn(res):
        count = 0
        final_results = []
    
        for idx, s in enumerate(tape_counts):
            # apply any batch transform post-processing
>           new_res = batch_fns[idx](res[count : count + s])
E           IndexError: too many indices for array: array is 0-dimensional, but 1 were indexed

/root/.pyenv/versions/3.8.16/lib/python3.8/site-packages/pennylane/transforms/batch_transform.py:510: IndexError
____________ TestDeviceIntegration.test_one_qubit_circuit[8192-d0] _____________
[gw1] linux -- Python 3.8.16 /root/.pyenv/versions/3.8.16/bin/python3.8

self = <test_integration.TestDeviceIntegration object at 0x7fcd904d4790>
shots = 8192
d = ('braket.aws.qubit', 'arn:aws:braket:::device/quantum-simulator/amazon/sv1')
tol = {'atol': 0.05, 'rtol': 0.1}
extra_kwargs = <function extra_kwargs.<locals>._extra_kwargs at 0x7fcd3f32a310>

    @pytest.mark.parametrize("d", shortname_and_backends)
    @pytest.mark.parametrize("shots", [None, 8192])
    def test_one_qubit_circuit(self, shots, d, tol, extra_kwargs):
        """Test that devices provide correct result for a simple circuit"""
        dev = TestDeviceIntegration._device(d, 1, extra_kwargs)
    
        a = 0.543
        b = 0.123
        c = 0.987
    
        @qml.qnode(dev)
        def circuit(x, y, z):
            """Reference QNode"""
            qml.BasisState(np.array([1]), wires=0)
            qml.Hadamard(wires=0)
            qml.Rot(x, y, z, wires=0)
            return qml.expval(qml.PauliZ(0))
    
>       assert np.allclose(circuit(a, b, c), np.cos(a) * np.sin(b), **tol)

test/integ_tests/test_integration.py:64: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 
/root/.pyenv/versions/3.8.16/lib/python3.8/site-packages/pennylane/qnode.py:867: in __call__
    res = qml.execute(
/root/.pyenv/versions/3.8.16/lib/python3.8/site-packages/pennylane/interfaces/execution.py:493: in execute
    return batch_fn(res)
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 

res = array(0.10504264)

    def processing_fn(res):
        count = 0
        final_results = []
    
        for idx, s in enumerate(tape_counts):
            # apply any batch transform post-processing
>           new_res = batch_fns[idx](res[count : count + s])
E           IndexError: too many indices for array: array is 0-dimensional, but 1 were indexed

/root/.pyenv/versions/3.8.16/lib/python3.8/site-packages/pennylane/transforms/batch_transform.py:510: IndexError
___________________ TestVar.test_var_hermitian[device0-None] ___________________
[gw0] linux -- Python 3.8.16 /root/.pyenv/versions/3.8.16/bin/python3.8

self = <test_var.TestVar object at 0x7f9b2e013b50>
device = <function device.<locals>._device at 0x7f9b1793bee0>, shots = None
tol = {'atol': 0.05, 'rtol': 0.1}

    def test_var_hermitian(self, device, shots, tol):
        """Tests for variance calculation using an arbitrary Hermitian observable"""
        dev = device(2)
    
        phi = 0.543
        theta = 0.6543
        H = np.array([[4, -1 + 6j], [-1 - 6j, 2]])
    
        @qml.qnode(dev)
        def circuit():
            qml.RX(phi, wires=[0])
            qml.RY(theta, wires=[0])
            return qml.var(qml.Hermitian(H, wires=[0]))
    
        expected = 0.5 * (
            2 * np.sin(2 * theta) * np.cos(phi) ** 2
            + 24 * np.sin(phi) * np.cos(phi) * (np.sin(theta) - np.cos(theta))
            + 35 * np.cos(2 * phi)
            + 39
        )
>       assert np.allclose(circuit(), expected, **tol)

test/integ_tests/test_var.py:63: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 
/root/.pyenv/versions/3.8.16/lib/python3.8/site-packages/pennylane/qnode.py:867: in __call__
    res = qml.execute(
/root/.pyenv/versions/3.8.16/lib/python3.8/site-packages/pennylane/interfaces/execution.py:493: in execute
    return batch_fn(res)
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 

res = array(27.38204933)

    def processing_fn(res):
        count = 0
        final_results = []
    
        for idx, s in enumerate(tape_counts):
            # apply any batch transform post-processing
>           new_res = batch_fns[idx](res[count : count + s])
E           IndexError: too many indices for array: array is 0-dimensional, but 1 were indexed

/root/.pyenv/versions/3.8.16/lib/python3.8/site-packages/pennylane/transforms/batch_transform.py:510: IndexError
________________________ TestVar.test_var[device0-None] ________________________
[gw1] linux -- Python 3.8.16 /root/.pyenv/versions/3.8.16/bin/python3.8

self = <test_var.TestVar object at 0x7fcd90462a30>
device = <function device.<locals>._device at 0x7fcd3f32a670>, shots = None
tol = {'atol': 0.05, 'rtol': 0.1}

    def test_var(self, device, shots, tol):
        """Tests for variance calculation"""
        dev = device(2)
    
        phi = 0.543
        theta = 0.6543
    
        @qml.qnode(dev)
        def circuit():
            qml.RX(phi, wires=[0])
            qml.RY(theta, wires=[0])
            return qml.var(qml.PauliZ(wires=[0]))
    
        expected = 0.25 * (3 - np.cos(2 * theta) - 2 * np.cos(theta) ** 2 * np.cos(2 * phi))
>       assert np.allclose(circuit(), expected, **tol)

test/integ_tests/test_var.py:41: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 
/root/.pyenv/versions/3.8.16/lib/python3.8/site-packages/pennylane/qnode.py:867: in __call__
    res = qml.execute(
/root/.pyenv/versions/3.8.16/lib/python3.8/site-packages/pennylane/interfaces/execution.py:493: in execute
    return batch_fn(res)
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 

res = array(0.53849344)

    def processing_fn(res):
        count = 0
        final_results = []
    
        for idx, s in enumerate(tape_counts):
            # apply any batch transform post-processing
>           new_res = batch_fns[idx](res[count : count + s])
E           IndexError: too many indices for array: array is 0-dimensional, but 1 were indexed

/root/.pyenv/versions/3.8.16/lib/python3.8/site-packages/pennylane/transforms/batch_transform.py:510: IndexError
________________ TestTensorVar.test_paulix_pauliy[device0-None] ________________
[gw0] linux -- Python 3.8.16 /root/.pyenv/versions/3.8.16/bin/python3.8

self = <test_var.TestTensorVar object at 0x7f9b2e013130>
device = <function device.<locals>._device at 0x7f9b2e1c7c10>, shots = None
tol = {'atol': 0.05, 'rtol': 0.1}

    def test_paulix_pauliy(self, device, shots, tol):
        """Test that a tensor product involving PauliX and PauliY works correctly"""
        dev = device(3)
    
        theta = 0.432
        phi = 0.123
        varphi = -0.543
    
        @qml.qnode(dev)
        def circuit():
            qml.RX(theta, wires=[0])
            qml.RX(phi, wires=[1])
            qml.RX(varphi, wires=[2])
            qml.CNOT(wires=[0, 1])
            qml.CNOT(wires=[1, 2])
            return qml.var(qml.PauliX(0) @ qml.PauliY(2))
    
        expected = (
            8 * np.sin(theta) ** 2 * np.cos(2 * varphi) * np.sin(phi) ** 2
            - np.cos(2 * (theta - phi))
            - np.cos(2 * (theta + phi))
            + 2 * np.cos(2 * theta)
            + 2 * np.cos(2 * phi)
            + 14
        ) / 16
>       assert np.allclose(circuit(), expected, **tol)

test/integ_tests/test_var.py:95: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 
/root/.pyenv/versions/3.8.16/lib/python3.8/site-packages/pennylane/qnode.py:867: in __call__
    res = qml.execute(
/root/.pyenv/versions/3.8.16/lib/python3.8/site-packages/pennylane/interfaces/execution.py:493: in execute
    return batch_fn(res)
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 

res = array(0.99929549)

    def processing_fn(res):
        count = 0
        final_results = []
    
        for idx, s in enumerate(tape_counts):
            # apply any batch transform post-processing
>           new_res = batch_fns[idx](res[count : count + s])
E           IndexError: too many indices for array: array is 0-dimensional, but 1 were indexed

/root/.pyenv/versions/3.8.16/lib/python3.8/site-packages/pennylane/transforms/batch_transform.py:510: IndexError
_______________ TestTensorVar.test_pauliz_hadamard[device0-None] _______________
[gw1] linux -- Python 3.8.16 /root/.pyenv/versions/3.8.16/bin/python3.8

self = <test_var.TestTensorVar object at 0x7fcd90517940>
device = <function device.<locals>._device at 0x7fcd3f32a280>, shots = None
tol = {'atol': 0.05, 'rtol': 0.1}

    def test_pauliz_hadamard(self, device, shots, tol):
        """Test that a tensor product involving PauliZ and PauliY and hadamard works correctly"""
        dev = device(3)
    
        theta = 0.432
        phi = 0.123
        varphi = -0.543
    
        @qml.qnode(dev)
        def circuit():
            qml.RX(theta, wires=[0])
            qml.RX(phi, wires=[1])
            qml.RX(varphi, wires=[2])
            qml.CNOT(wires=[0, 1])
            qml.CNOT(wires=[1, 2])
            return qml.var(qml.PauliZ(wires=[0]) @ qml.Hadamard(wires=[1]) @ qml.PauliY(wires=[2]))
    
        expected = (
            3
            + np.cos(2 * phi) * np.cos(varphi) ** 2
            - np.cos(2 * theta) * np.sin(varphi) ** 2
            - 2 * np.cos(theta) * np.sin(phi) * np.sin(2 * varphi)
        ) / 4
    
>       assert np.allclose(circuit(), expected, **tol)

test/integ_tests/test_var.py:121: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 
/root/.pyenv/versions/3.8.16/lib/python3.8/site-packages/pennylane/qnode.py:867: in __call__
    res = qml.execute(
/root/.pyenv/versions/3.8.16/lib/python3.8/site-packages/pennylane/interfaces/execution.py:493: in execute
    return batch_fn(res)
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 

res = array(0.93368122)

    def processing_fn(res):
        count = 0
        final_results = []
    
        for idx, s in enumerate(tape_counts):
            # apply any batch transform post-processing
>           new_res = batch_fns[idx](res[count : count + s])
E           IndexError: too many indices for array: array is 0-dimensional, but 1 were indexed

/root/.pyenv/versions/3.8.16/lib/python3.8/site-packages/pennylane/transforms/batch_transform.py:510: IndexError
__________________ TestTensorVar.test_hermitian[device0-None] __________________
[gw0] linux -- Python 3.8.16 /root/.pyenv/versions/3.8.16/bin/python3.8

self = <test_var.TestTensorVar object at 0x7f9b2e0133d0>
device = <function device.<locals>._device at 0x7f9b144a69d0>, shots = None
tol = {'atol': 0.05, 'rtol': 0.1}

    def test_hermitian(self, device, shots, tol):
        """Test that a tensor product involving qml.Hermitian works correctly"""
        dev = device(3)
    
        theta = 0.432
        phi = 0.123
        varphi = -0.543
        A = np.array(
            [
                [-6, 2 + 1j, -3, -5 + 2j],
                [2 - 1j, 0, 2 - 1j, -5 + 4j],
                [-3, 2 + 1j, 0, -4 + 3j],
                [-5 - 2j, -5 - 4j, -4 - 3j, -6],
            ]
        )
    
        @qml.qnode(dev)
        def circuit():
            qml.RX(theta, wires=[0])
            qml.RX(phi, wires=[1])
            qml.RX(varphi, wires=[2])
            qml.CNOT(wires=[0, 1])
            qml.CNOT(wires=[1, 2])
            return qml.var(qml.PauliZ(wires=[0]) @ qml.Hermitian(A, wires=[1, 2]))
    
        expected = (
            1057
            - np.cos(2 * phi)
            + 12 * (27 + np.cos(2 * phi)) * np.cos(varphi)
            - 2 * np.cos(2 * varphi) * np.sin(phi) * (16 * np.cos(phi) + 21 * np.sin(phi))
            + 16 * np.sin(2 * phi)
            - 8 * (-17 + np.cos(2 * phi) + 2 * np.sin(2 * phi)) * np.sin(varphi)
            - 8 * np.cos(2 * theta) * (3 + 3 * np.cos(varphi) + np.sin(varphi)) ** 2
            - 24 * np.cos(phi) * (np.cos(phi) + 2 * np.sin(phi)) * np.sin(2 * varphi)
            - 8
            * np.cos(theta)
            * (
                4
                * np.cos(phi)
                * (
                    4
                    + 8 * np.cos(varphi)
                    + np.cos(2 * varphi)
                    - (1 + 6 * np.cos(varphi)) * np.sin(varphi)
                )
                + np.sin(phi)
                * (
                    15
                    + 8 * np.cos(varphi)
                    - 11 * np.cos(2 * varphi)
                    + 42 * np.sin(varphi)
                    + 3 * np.sin(2 * varphi)
                )
            )
        ) / 16
>       assert np.allclose(circuit(), expected, **tol)

test/integ_tests/test_var.py:178: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 
/root/.pyenv/versions/3.8.16/lib/python3.8/site-packages/pennylane/qnode.py:867: in __call__
    res = qml.execute(
/root/.pyenv/versions/3.8.16/lib/python3.8/site-packages/pennylane/interfaces/execution.py:493: in execute
    return batch_fn(res)
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 

res = array(47.71487867)

    def processing_fn(res):
        count = 0
        final_results = []
    
        for idx, s in enumerate(tape_counts):
            # apply any batch transform post-processing
>           new_res = batch_fns[idx](res[count : count + s])
E           IndexError: too many indices for array: array is 0-dimensional, but 1 were indexed

/root/.pyenv/versions/3.8.16/lib/python3.8/site-packages/pennylane/transforms/batch_transform.py:510: IndexError
=============================== warnings summary ===============================
test/integ_tests/test_apply.py: 45 warnings
test/integ_tests/test_adjoint_gradient.py: 1 warning
test/integ_tests/test_expval.py: 23 warnings
test/integ_tests/test_integration.py: 2 warnings
test/integ_tests/test_sample.py: 6 warnings
test/integ_tests/test_var.py: 10 warnings
test/integ_tests/test_tracking.py: 1 warning
  /root/.pyenv/versions/3.8.16/lib/python3.8/site-packages/braket/default_simulator/simulator.py:265: UserWarning: You are running a noise-free circuit on the density matrix simulator. Consider running this circuit on the state vector simulator: LocalSimulator("default") for a better user experience.
    warnings.warn(

test/integ_tests/test_adjoint_gradient.py::TestAdjointGradient::test_grad[on_demand_sv_device0-None]
test/integ_tests/test_adjoint_gradient.py::TestAdjointGradient::test_grad_large_circuit[on_demand_sv_device0-None]
test/integ_tests/test_adjoint_gradient.py::TestAdjointGradient::test_default_diff_method[device0-None]
  /codebuild/output/src291667968/src/src/braket/pennylane_plugin/braket_device.py:239: VisibleDeprecationWarning: Creating an ndarray from ragged nested sequences (which is a list-or-tuple of lists-or-tuples-or ndarrays with different lengths or shapes) is deprecated. If you meant to do this, you must specify 'dtype=object' when creating the ndarray.
    return onp.array(results).squeeze()

test/integ_tests/test_ahs_device.py::TestQnodeIntegration::test_circuit_can_be_called_global_drive[H5-params5]
test/integ_tests/test_ahs_device.py::TestQnodeIntegration::test_circuit_can_be_called_local_detunings[amp-local_params0-local_wires0-H5-params5]
test/integ_tests/test_ahs_device.py::TestQnodeIntegration::test_circuit_can_be_called_local_detunings[4.5-local_params1-local_wires1-H5-params5]
  /root/.pyenv/versions/3.8.16/lib/python3.8/site-packages/braket/analog_hamiltonian_simulator/rydberg/validators/field_validator_util.py:22: UserWarning: Value 0 (25132741.228718344) in amplitude time series outside the typical range [0, 25000000.0]. The values should  be specified in SI units.
    warnings.warn(

test/integ_tests/test_batch.py::test_batch_execution_of_gradient_torch[device0-None]
test/integ_tests/test_batch.py::test_batch_execution_of_gradient_torch[device0-None]
  /root/.pyenv/versions/3.8.16/lib/python3.8/site-packages/semantic_version/base.py:600: PendingDeprecationWarning: The Spec() class will be removed in 3.1; use SimpleSpec() instead.
    return Spec(spec).match(Version(version))

-- Docs: https://docs.pytest.org/en/stable/how-to/capture-warnings.html
- generated xml file: /codebuild/output/src291667968/src/test-reports/junit.xml -
=========================== short test summary info ============================
FAILED test/integ_tests/test_expval.py::TestExpval::test_paulix_expectation[device0-None]
FAILED test/integ_tests/test_expval.py::TestExpval::test_pauliy_expectation[device0-None]
FAILED test/integ_tests/test_expval.py::TestExpval::test_hadamard_expectation[device0-None]
FAILED test/integ_tests/test_expval.py::TestExpval::test_hermitian_expectation[device0-None]
FAILED test/integ_tests/test_expval.py::TestExpval::test_hamiltonian[device0-None]
FAILED test/integ_tests/test_expval.py::TestExpval::test_nondiff_param[device0-None]
FAILED test/integ_tests/test_expval.py::TestTensorExpval::test_paulix_pauliy[device0-None]
FAILED test/integ_tests/test_expval.py::TestTensorExpval::test_pauliz_hadamard[device0-None]
FAILED test/integ_tests/test_expval.py::TestExpval::test_identity_expectation[device0-None]
FAILED test/integ_tests/test_expval.py::TestExpval::test_pauliz_expectation[device0-None]
FAILED test/integ_tests/test_integration.py::TestDeviceIntegration::test_one_qubit_circuit[None-d0]
FAILED test/integ_tests/test_integration.py::TestDeviceIntegration::test_one_qubit_circuit[8192-d0]
FAILED test/integ_tests/test_var.py::TestVar::test_var_hermitian[device0-None]
FAILED test/integ_tests/test_var.py::TestVar::test_var[device0-None] - IndexE...
FAILED test/integ_tests/test_var.py::TestTensorVar::test_paulix_pauliy[device0-None]
FAILED test/integ_tests/test_var.py::TestTensorVar::test_pauliz_hadamard[device0-None]
FAILED test/integ_tests/test_var.py::TestTensorVar::test_hermitian[device0-None]
====== 17 failed, 312 passed, 9 skipped, 96 warnings in 325.66s (0:05:25) ======