patch qml.matrix to accept qnodes compiled by catalyst.qjit (#5753)

**Context:** Currently `qml.matrix`
 breaks when the input is a catalyst-compiled QNode instead of a plain
QNode, and the QNode need to be manually retrieved by the user: 
PennyLaneAI/catalyst#765

It would be beneficial to do this "unwrap" in `qml.matrix` , like those
in qml.draw:

https://github.com/PennyLaneAI/pennylane/blob/59a1e0586e707d057a0c92d4239036afa5312b73/pennylane/drawer/draw.py#L213-L214

**Description of the Change:** 
1. in `pennylane/ops/functions/matrix.py`, in `qml.matrix(op)`, if the
input `op` is a catalyst.qjit compiled function, dispatches the behavior
to be `qml.matrix(op.user_function)`.
2. Added a test in `tests/ops/functions/test_matrix.py`

**Benefits:** a qjit compiled qnode can be passed into `qml.matrix`
directly to query the matrix representation of the circuit


**Related GitHub Issues:**
PennyLaneAI/catalyst#765

[sc-64247]

doc/releases/changelog-dev.md

@@ -239,6 +239,9 @@
 * Fixes a bug in `qml.math.dot` that raises an error when only one of the operands is a scalar.
   [(#5702)](https://github.com/PennyLaneAI/pennylane/pull/5702)
 
+* `qml.matrix` is now compatible with qnodes compiled by catalyst.qjit.
+  [(#5753)](https://github.com/PennyLaneAI/pennylane/pull/5753)
+
 <h3>Contributors ✍️</h3>
 
 This release contains contributions from (in alphabetical order):
@@ -259,4 +262,5 @@ Vincent Michaud-Rioux,
 Lee James O'Riordan,
 Mudit Pandey,
 Kenya Sakka,
+Haochen Paul Wang,
 David Wierichs.

pennylane/drawer/draw.py

@@ -18,7 +18,6 @@
 """
 import warnings
 from functools import wraps
-from importlib.metadata import distribution
 
 import pennylane as qml
 
@@ -27,12 +26,8 @@
 
 
 def catalyst_qjit(qnode):
-    """The ``catalyst.while`` wrapper method"""
-    try:
-        distribution("pennylane_catalyst")
-        return qnode.__class__.__name__ == "QJIT"
-    except ImportError:
-        return False
+    """A method checking whether a qnode is compiled by catalyst.qjit"""
+    return qnode.__class__.__name__ == "QJIT" and hasattr(qnode, "user_function")
 
 
 def draw(

pennylane/ops/functions/matrix.py

@@ -27,6 +27,11 @@
 from pennylane.typing import TensorLike
 
 
+def catalyst_qjit(qnode):
+    """A method checking whether a qnode is compiled by catalyst.qjit"""
+    return qnode.__class__.__name__ == "QJIT" and hasattr(qnode, "user_function")
+
+
 def matrix(op: Union[Operator, PauliWord, PauliSentence], wire_order=None) -> TensorLike:
     r"""The matrix representation of an operation or quantum circuit.
 
@@ -177,6 +182,9 @@ def circuit():
             wires specified, and this is the order in which wires appear in ``circuit()``.
 
     """
+    if catalyst_qjit(op):
+        op = op.user_function
+
     if not isinstance(op, Operator):
 
         if isinstance(op, (PauliWord, PauliSentence)):

tests/ops/functions/test_matrix.py

@@ -659,6 +659,34 @@ def circuit(theta):
 
         assert np.allclose(matrix, expected_matrix)
 
+    @pytest.mark.catalyst
+    @pytest.mark.external
+    def test_catalyst(self):
+        """Test with Catalyst interface"""
+
+        import catalyst
+
+        dev = qml.device("lightning.qubit", wires=1)
+
+        # create a plain QNode
+        @qml.qnode(dev)
+        def f():
+            qml.PauliX(0)
+            return qml.state()
+
+        # create a qjit-compiled QNode by decorating a function
+        @catalyst.qjit
+        @qml.qnode(dev)
+        def g():
+            qml.PauliX(0)
+            return qml.state()
+
+        # create a qjit-compiled QNode by passing in the plain QNode directly
+        h = catalyst.qjit(f)
+
+        assert np.allclose(f(), g(), h())
+        assert np.allclose(qml.matrix(f)(), qml.matrix(g)(), qml.matrix(h)())
+
     @pytest.mark.jax
     def test_get_unitary_matrix_interface_jax(self):
         """Test with JAX interface"""

tests/pytest.ini

@@ -18,6 +18,7 @@ markers =
     param-shift: marks tests for the parameter shift (deselect with '-m "not param-shift"')
     logging: marks tests for pennylane logging
     external: marks tests that require external packages such as matplotlib and PyZX
+    catalyst: marks tests for catalyst testing (select with '-m "catalyst"')
 filterwarnings =
     ignore::DeprecationWarning:autograd.numpy.numpy_wrapper
     ignore:Casting complex values to real::autograd.numpy.numpy_wrapper