Extend Partial trace function support for PennyLane operators

[babcockt18]

Feature details

As the title suggests, this feature would extend qml.math.partial_trace functionality to handle PennyLane operators. The public facing function already exists. This issue is part of an optional enhancement that was proposed by @Qottmann but was closed before the latter portion of the issue was implemented.

Implementation

• [ ] The function qml.math.partial_trace supports PennyLane operators. E.g.

>>> qml.math.partial_trace(op, wires=[1])
X(0) + Y(0)
>>> qml.math.partial_trace(op, wires=[0])
X(1) + Y(1)

• [ ] updated docstring to include descriptive doc string that define inputs, outputs and provides at least one illustrative example specifically for an operator
• [ ] A changelog entry is added to doc/releases/changelog-dev.md
• [ ] Unit tests are added to tests/math/test_matrix_manipulation.py ensuring inputs and outputs of the function are correct (in particular checking the supported dtypes float and complex and the supported interfaces jax, torch, autograd, tensorflow)
• [ ] Integration tests are added to tests/math/test_matrix_manipulation.py and in particular should cover instances of SProd, Sum and Prod classes (pennylane operator classes)

How important would you say this feature is?

1: Not important. Would be nice to have.

Additional information

No response

[trbromley]

Thanks @babcockt18! Yes this sounds like a nice idea and should be achievable for sums/products of the standard observables like X, Y, Z, etc. We could initially restrict to operators that are Pauli sentences (linear combinations of tensor products of Paulis) to make things easier.

Also, this partial_trace doesn't need to be the same one as in the math module. We could place it in op_math and also consider importing at top level.

Is this something you'd be interested to take on?

[babcockt18]

@trbromley Awesome! definitely want to take this on! Thanks for asking! if you could assign it to me that would be great!

[trbromley]

Sorry @babcockt18, I lost track of replying to this. Would you still be interested on working on this feature?

[babcockt18]

@trbromley All good! Still interested!

[Alex-Preciado]

Great! @babcockt18, please don't hesitate to reach out if you have any questions. We'll be happy to assist with the code review once a pull request is ready.