Add a QutritChannel operation

[trbromley]

Feature details

Version 0.36 of PennyLane includes a "default.qutrit.mixed" device that is capable of simulating qutrit systems in the mixed-state setting. This unlocks the ability to simulate qutrit-based noise, and we would like to add more qutrit noise channels to complement our existing range of qutrit operators.

This objective of this issue is to provide a qutrit analogue of QubitChannel, which would allow noise to be specified via a collection of $(3\times 3)$-dimensional Kraus operators.

Implementation

It should be possible to execute the following code:

import pennylane as qml
import numpy as np

kraus = [  # See https://github.com/PennyLaneAI/pennylane/pull/5503
    np.array([[1, 0, 0],
              [0, 0.70710678, 0],
              [0, 0, 0.8660254]]),
    np.array([[0, 0.70710678, 0],
              [0, 0, 0],
              [0, 0, 0]]),
    np.array([[0, 0, 0.5],
              [0, 0, 0],
              [0, 0, 0]])
]

dev = qml.device("default.qutrit.mixed")

@qml.qnode(dev)
def f():
    qml.QutritChannel(kraus, 0)
    return qml.expval(qml.GellMann(wires=0, index=1))

f()

This can be done by following these instructions:

1. Implement a qml.QutritChannel operation in pennylane/ops/channel.py, which defines the qutrit equivalent of the qml.QubitChannel operation. This operation should:

   • Take a list containing three $3 \times 3$ Kraus matrices, along with wires and id (optional string representation of the operator).
   • Define the compute_kraus_matrices, which is a staticmethod that returns the Kraus matrices in a list, and _flatten functions (just like in qml.QubitChannel).

2. Add unit and integration tests:

   • Unit tests must verify the correctness of the operation, its attributes, and its methods.
   • Integration tests must show that the output of a QNode that applies the qml.QutritChannel operation is correct.
   • Integration tests must also verify that QNodes that apply the qml.QutritChannel operation can be differentiated using the parameter-shift, finite-diff, and backprop differentiation methods and that the computed gradients/jacobians are correct.
   • Integration tests must be implemented for both the default.qutrit and default.qutrit.mixed devices.

A successful PR must include:

• The requested feature
• Comprehensive tests
• Relevant updates to the documentation
• An entry in the release notes
• Passing all CI tests

Check out our development guide for more information, and don't hesitate to reach out to us if you need assistance.

How important would you say this feature is?

1: Not important. Would be nice to have.

Additional information

Existing work is underway to support specific qutrit channels in PRs #5502 and #5503.

[Tarun-Kumar07]

Hi @trbromley , I want to tackle this issue as part of unitaryhack.
Could you please assign it to me?

[trbromley]

Nice @Tarun-Kumar07! Let us know how we can help as you get started.

Regarding assigning, we typically don't assign the issue until a resolving PR has fixed it. We don't want to discourage other attempts at the bounty, and it's also encouraged if contributors want to optionally work together.

[Tarun-Kumar07]

No problem, thanks @trbromley

[tnemoz]

I'm a bit puzzled by the existing QubitChannel class.

Currently, QubitChannel works like an arbitrary channel represented by Kraus operators: the checks therein never check for the dimension of the matrices to be $2\times2$. So, adding a QutritChannel the same way in pennylane/ops/channel.py would just result in code duplication between the two.

At some point I thought "Ok, but then one can just add the checks that the matrices are $2\times2$, and similarly for the qutrit case", but these lines:

# check the dimension of all Kraus matrices are valid
if any(K.ndim != 2 for K in K_list):
    raise ValueError(
        "Dimension of all Kraus matrices must be (2**num_wires, 2**num_wires)."
    )

show that it is intended to accept Kraus operators in higher dimensions.

So, is it preferable to:

1. Create an abstract KrausChannel class that implements these generic checks, and make both QubitChannel and QutritChannel inherit from it?
2. Same as 1., but additionally check that the dimensions matches what we expect to see?
3. Copy QubitChannel into QutritChannel but add the check that the dimensions are expected to be $2\times2$ and $3\times3$?
4. Copy without these additional checks?

For instance, for now, nothing prevents to instantiate QubitChannel with a list of $3\times3$ matrices. Of course, then the decomposer yells at us, but my question is just whether this is the expected behavior or not.

---

On an unrelated topic, I wanted to test a similar code of yours to see how QubitChannel works:

import pennylane as qml
import numpy as np

kraus_X = [
    np.array([
        [0., 1.],
        [1., 0.]
    ])
]

dev = qml.device("default.qubit")

@qml.qnode(dev)
def f():
    #qml.QubitChannel(kraus_X, 0)
    qml.X(wires=0)
    return qml.expval(qml.PauliZ(wires=0))

f()

Running it with qml.X works flawlessly, but using qml.QubitChannel returns this error:

pennylane._device.DeviceError: Operator QubitChannel(array([[0., 1.],
       [1., 0.]]), wires=[0]) not supported with default.qubit and does not provide a decomposition.

Did I misunderstand something?

Sorry for the long comment!

[Tarun-Kumar07]

Hi @trbromley , I have a doubt.
Do qutrit channels support the device default.qutrit ?
Because the test of QutritDepolarizingChannel fails when I change the device to default.qutrit. The modified test

    @pytest.mark.parametrize("angle", np.linspace(0, 2 * np.pi, 7))
    def test_grad_depolarizing(self, angle):
        """Test that analytical gradient is computed correctly for different states. Channel
        grad recipes are independent of channel parameter"""

        dev = qml.device("default.qutrit", wires=1)
        prob = pnp.array(0.5, requires_grad=True)

        @qml.qnode(dev, diff_method="parameter-shift")
        def circuit(p):
            qml.TRX(angle, wires=0, subspace=(0, 1))
            qml.TRX(angle, wires=0, subspace=(1, 2))
            qml.QutritDepolarizingChannel(p, wires=0)
            return qml.expval(qml.GellMann(0, 3) + qml.GellMann(0, 8))

        expected_errorless = (
            (np.sqrt(3) - 3) * (1 - np.cos(2 * angle)) / 24
            - 2 / np.sqrt(3) * np.sin(angle / 2) ** 4
            + (np.sqrt(1 / 3) + 1) * np.cos(angle / 2) ** 2
        )

        assert np.allclose(circuit(prob), ((prob - (1 / 9)) / (8 / 9)) * expected_errorless)

        gradient = np.squeeze(qml.grad(circuit)(prob))
        assert np.allclose(gradient, circuit(1) - circuit(0))
        assert np.allclose(gradient, -(9 / 8) * expected_errorless)

When I add QutritDepolarizingChannel to operations of default.qutrit it fails with the stack trace.

test_qutrit_channel_ops.py:91: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 
../../../pennylane/workflow/qnode.py:1098: in __call__
    res = self._execution_component(args, kwargs, override_shots=override_shots)
../../../pennylane/workflow/qnode.py:1052: in _execution_component
    res = qml.execute(
../../../pennylane/workflow/execution.py:784: in execute
    results = ml_boundary_execute(tapes, execute_fn, jpc, device=device)
../../../pennylane/workflow/interfaces/autograd.py:147: in autograd_execute
    return _execute(parameters, tuple(tapes), execute_fn, jpc)
/opt/homebrew/Caskroom/miniconda/base/envs/pennylane-dev/lib/python3.11/site-packages/autograd/tracer.py:48: in f_wrapped
    return f_raw(*args, **kwargs)
../../../pennylane/workflow/interfaces/autograd.py:168: in _execute
    return execute_fn(tapes)
../../../pennylane/workflow/execution.py:298: in inner_execute
    results = device_execution(transformed_tapes)
/opt/homebrew/Caskroom/miniconda/base/envs/pennylane-dev/lib/python3.11/contextlib.py:81: in inner
    return func(*args, **kwds)
../../../pennylane/_qubit_device.py:501: in batch_execute
    res = self.execute(circuit)
../../../pennylane/_qubit_device.py:292: in execute
    self.apply(
../../../pennylane/devices/default_qutrit.py:187: in apply
    self._state = self._apply_operation(self._state, operation)
../../../pennylane/devices/default_qutrit.py:249: in _apply_operation
    matrix = self._asarray(self._get_unitary_matrix(operation), dtype=self.C_DTYPE)
../../../pennylane/devices/default_qutrit.py:370: in _get_unitary_matrix
    return unitary.matrix()
../../../pennylane/operation.py:838: in matrix
    canonical_matrix = self.compute_matrix(*self.parameters, **self.hyperparameters)
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 

params = (0.5,), hyperparams = {}

    @staticmethod
    def compute_matrix(*params, **hyperparams) -> TensorLike:  # pylint:disable=unused-argument
        r"""Representation of the operator as a canonical matrix in the computational basis (static method).
    
        The canonical matrix is the textbook matrix representation that does not consider wires.
        Implicitly, this assumes that the wires of the operator correspond to the global wire order.
    
        .. seealso:: :meth:`.Operator.matrix` and :func:`qml.matrix() <pennylane.matrix>`
    
        Args:
            *params (list): trainable parameters of the operator, as stored in the ``parameters`` attribute
            **hyperparams (dict): non-trainable hyperparameters of the operator, as stored in the ``hyperparameters`` attribute
    
        Returns:
            tensor_like: matrix representation
        """
>       raise MatrixUndefinedError
E       pennylane.operation.MatrixUndefinedError

[Tarun-Kumar07]

@tnemoz, I think the below behaviour is wrong

For instance, for now, nothing prevents to instantiate QubitChannel with a list of 
 matrices. Of course, then the decomposer yells at us, but my question is just whether this is the expected behavior or not.

Reason being for QubitChannel kraus matrices should be (2,2) assuming num_wires=1. I think the validation condition should have been

if any(K.shape[0] != 2 for K in K_list):

because the condition right now

if any(K.ndim != 2 for K in K_list):

will return true for both matrices of size (3x3) and (2x2) as number of dimensions is 2.

[tnemoz]

@Tarun-Kumar07 I totally agree, I just wanted a confirmation that this wasn't intentional before going on!

[mudit2812]

Hi @trbromley , I have a doubt. Do qutrit channels support the device default.qutrit ? Because the test of QutritDepolarizingChannel fails when I change the device to default.qutrit. The modified test

    @pytest.mark.parametrize("angle", np.linspace(0, 2 * np.pi, 7))
    def test_grad_depolarizing(self, angle):
        """Test that analytical gradient is computed correctly for different states. Channel
        grad recipes are independent of channel parameter"""

        dev = qml.device("default.qutrit", wires=1)
        prob = pnp.array(0.5, requires_grad=True)

        @qml.qnode(dev, diff_method="parameter-shift")
        def circuit(p):
            qml.TRX(angle, wires=0, subspace=(0, 1))
            qml.TRX(angle, wires=0, subspace=(1, 2))
            qml.QutritDepolarizingChannel(p, wires=0)
            return qml.expval(qml.GellMann(0, 3) + qml.GellMann(0, 8))

        expected_errorless = (
            (np.sqrt(3) - 3) * (1 - np.cos(2 * angle)) / 24
            - 2 / np.sqrt(3) * np.sin(angle / 2) ** 4
            + (np.sqrt(1 / 3) + 1) * np.cos(angle / 2) ** 2
        )

        assert np.allclose(circuit(prob), ((prob - (1 / 9)) / (8 / 9)) * expected_errorless)

        gradient = np.squeeze(qml.grad(circuit)(prob))
        assert np.allclose(gradient, circuit(1) - circuit(0))
        assert np.allclose(gradient, -(9 / 8) * expected_errorless)

When I add QutritDepolarizingChannel to operations of default.qutrit it fails with the stack trace.

test_qutrit_channel_ops.py:91: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 
../../../pennylane/workflow/qnode.py:1098: in __call__
    res = self._execution_component(args, kwargs, override_shots=override_shots)
../../../pennylane/workflow/qnode.py:1052: in _execution_component
    res = qml.execute(
../../../pennylane/workflow/execution.py:784: in execute
    results = ml_boundary_execute(tapes, execute_fn, jpc, device=device)
../../../pennylane/workflow/interfaces/autograd.py:147: in autograd_execute
    return _execute(parameters, tuple(tapes), execute_fn, jpc)
/opt/homebrew/Caskroom/miniconda/base/envs/pennylane-dev/lib/python3.11/site-packages/autograd/tracer.py:48: in f_wrapped
    return f_raw(*args, **kwargs)
../../../pennylane/workflow/interfaces/autograd.py:168: in _execute
    return execute_fn(tapes)
../../../pennylane/workflow/execution.py:298: in inner_execute
    results = device_execution(transformed_tapes)
/opt/homebrew/Caskroom/miniconda/base/envs/pennylane-dev/lib/python3.11/contextlib.py:81: in inner
    return func(*args, **kwds)
../../../pennylane/_qubit_device.py:501: in batch_execute
    res = self.execute(circuit)
../../../pennylane/_qubit_device.py:292: in execute
    self.apply(
../../../pennylane/devices/default_qutrit.py:187: in apply
    self._state = self._apply_operation(self._state, operation)
../../../pennylane/devices/default_qutrit.py:249: in _apply_operation
    matrix = self._asarray(self._get_unitary_matrix(operation), dtype=self.C_DTYPE)
../../../pennylane/devices/default_qutrit.py:370: in _get_unitary_matrix
    return unitary.matrix()
../../../pennylane/operation.py:838: in matrix
    canonical_matrix = self.compute_matrix(*self.parameters, **self.hyperparameters)
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 

params = (0.5,), hyperparams = {}

    @staticmethod
    def compute_matrix(*params, **hyperparams) -> TensorLike:  # pylint:disable=unused-argument
        r"""Representation of the operator as a canonical matrix in the computational basis (static method).
    
        The canonical matrix is the textbook matrix representation that does not consider wires.
        Implicitly, this assumes that the wires of the operator correspond to the global wire order.
    
        .. seealso:: :meth:`.Operator.matrix` and :func:`qml.matrix() <pennylane.matrix>`
    
        Args:
            *params (list): trainable parameters of the operator, as stored in the ``parameters`` attribute
            **hyperparams (dict): non-trainable hyperparameters of the operator, as stored in the ``hyperparameters`` attribute
    
        Returns:
            tensor_like: matrix representation
        """
>       raise MatrixUndefinedError
E       pennylane.operation.MatrixUndefinedError

Hi @Tarun-Kumar07 , this is expected. default.qutrit is a pure state simulation device, and channels are only usable on mixed-state devices. In the qutrit regime, the only one currently available is default.qutrit.mixed.

@tnemoz , this is the same reason why you are unable to use QubitChannel with default.qubit, that that is a pure state device, while channels are meant to only be useable with mixed state devices (default.mixed).

[mudit2812]

@tnemoz @Tarun-Kumar07

So, is it preferable to:

1. Create an abstract KrausChannel class that implements these generic checks, and make both QubitChannel and QutritChannel inherit from it?
2. Same as 1., but additionally check that the dimensions matches what we expect to see?
3. Copy QubitChannel into QutritChannel but add the check that the dimensions are expected to be $2\times 2$ and $3\times 3$?
4. Copy without these additional checks?

I think that either 1 or 3 are reasonable solutions. While I do agree that it is generally better practice to use a parent class for cases like this with a lot of duplicate code, QubitChannel has only three reasonably simple methods. Additionally, such abstraction would be more helpful if we were to add QuditChannel or something else of the sort where we would want to generalize the dimensions, but we do not have any plans for that in our roadmap at the moment.

With that in mind, I would leave it to your discretion to choose either option 1 or 3 as I think that in this specific situation they're both valid.

[MashAliK]

@mudit2812

Hello,

I went with approach 1 and I'm at the testing stage but I'm having some trouble with the differentiation part. I tried following a method similar to test_grad_thermal_relaxation_error, so I have the following:

dev = qml.device("default.qutrit.mixed", wires=1)
p = np.array(0.5, requires_grad=True)

@qml.qnode(dev)
def circuit(param):
    K_list1 = [  
        np.sqrt(1-param)*np.array([[1, 0, 0],
                [0, 1, 0],
                [0, 0, 1]]),
        np.sqrt(param)*np.array([[1, 0, 0],
                [0, 0, 0],
                [0, 0, 0]]),
        np.sqrt(param)*np.array([[0, 0, 0],
                [0, 1, 0],
                [0, 0, 0]]),
        np.sqrt(param)*np.array([[0, 0, 0],
                [0, 0, 0],
                [0, 0, 1]])
    ]
    qml.QutritChannel(K_list1, 0)
    return qml.expval(qml.GellMann(wires=0, index=1))

print(qml.grad(circuit)(p))

I get the following error when I run this

...
File "/../pennylane/ops/channel.py", line 736, in __init__
    Kraus_sum = np.einsum("ajk,ajl->kl", K_arr.conj(), K_arr)
                                         ^^^^^^^^^^^^
TypeError: loop of ufunc does not support argument 0 of type ArrayBox which has no callable conjugate method

But this is a line that was originally in QubitChannel. I encounter the same error when running a similar test with QubitChannel. Could you let me know if I'm approaching this the wrong way? Thank you.

[trbromley]

Hi everyone! It's great to see some engagement around this bounty 🤩

I just to clarify how things work in case of multiple contribution attempts for this bounty:

• We'll merge the first complete PR that solves the issue, which should include necessary updates to tests and documentation. The author(s) of that issue will jointly win the bounty and share the amount.
• It is possible to work together on a PR, but all of the PR authors should be on board with doing so.
• Even if a PR is already open, it might still be worthwhile to make your own attempt in case there are any gaps in the already-open PR.

So far, @Tarun-Kumar07 has opened a PR and we'll be taking a look at that soon.

Thanks!

[mudit2812]

Hi @MashAliK . QubitChannel is not differentiable, so it is not an expectation for QutritChannel to be differentiable either. Note that this does not mean that you can't use QutritChannel in differentiable circuits, just that QutritChannel cannot have differentiable parameters.

[Tarun-Kumar07]

Hi @trbromley and @mudit2812 ,
Can this issue be assigned to me, as #5793 is merged.

[trbromley]

Awesome job @Tarun-Kumar07! Congratulations on completing this bounty 🚀

[trbromley]

@Tarun-Kumar07, would you optionally like to be tagged in any of our announcements around unitaryHACK or the upcoming PennyLane 0.37 release at the start of July? This would be on LinkedIn and Twitter, so if you would like, please share your username(s).

[Tarun-Kumar07]

Hey @trbromley , I would love that :).
My socials

• LinkedIn profile : Tarun Kumar Allamsetty
• Twitter : Tarun_Kumar_A