Repository: PennyLaneAI/pennylane (python)

Issue:
[BUG] `BasisRotation` is not JAX compatible
### Expected behavior

`qml.BasisRotation` is not supported with `jax.jit` or `qml.qjit`.

### Actual behavior

```python
dev = qml.device("lightning.qubit", wires=4)

@jax.jit
@qml.qnode(dev)
def f(U):
    for i in range(4):
        qml.Hadamard(i)
    qml.BasisRotation(unitary_matrix=U, wires=[0, 1, 2, 3], check=False)
    return qml.expval(qml.PauliZ(0))
```

```pycon
>>> from scipy.stats import unitary_group
>>> U = unitary_group.rvs(4)
>>> f(U)
[/usr/local/lib/python3.10/dist-packages/pennylane/operation.py](https://localhost:8080/#) in decomposition(self)
   1318             list[Operator]: decomposition of the operator
   1319         """
-> 1320         return self.compute_decomposition(
   1321             *self.parameters, wires=self.wires, **self.hyperparameters
   1322         )

[/usr/local/lib/python3.10/dist-packages/pennylane/templates/subroutines/basis_rotation.py](https://localhost:8080/#) in compute_decomposition(wires, unitary_matrix, check)
    169 
    170         op_list = []
--> 171         phase_list, givens_list = givens_decomposition(unitary_matrix)
    172 
    173         for idx, phase in enumerate(phase_list):

[/usr/local/lib/python3.10/dist-packages/pennylane/qchem/givens_decomposition.py](https://localhost:8080/#) in givens_decomposition(unitary)
    149     """
    150 
--> 151     unitary, (M, N) = qml.math.toarray(unitary).copy(), unitary.shape
    152     if M != N:
    153         raise ValueError(f"The unitary matrix should be of shape NxN, got {unitary.shape}")

[/usr/local/lib/python3.10/dist-packages/autoray/autoray.py](https://localhost:8080/#) in do(fn, like, *args, **kwargs)
     79     backend = _choose_backend(fn, args, kwargs, like=like)
     80     func = get_lib_fn(backend, fn)
---> 81     return func(*args, **kwargs)
     82 
     83 

[/usr/local/lib/python3.10/dist-packages/pennylane/math/single_dispatch.py](https://localhost:8080/#) in _to_numpy_jax(x)
    783         return np.array(getattr(x, "val", x))
    784     except TracerArrayConversionError as e:
--> 785         raise ValueError(
    786             "Converting a JAX array to a NumPy array not supported when using the JAX JIT."
    787         ) from e

ValueError: Converting a JAX array to a NumPy array not supported when using the JAX JIT.
```

### Additional information

This is occurring because the decomposition for `qml.BasisRotation` calls [`qml.qchem.givens_decomposition`](https://github.com/PennyLaneAI/pennylane/blob/3bae4a2c2ea69e7734be8fbeb7641ae7cfe51a52/pennylane/qchem/givens_decomposition.py#L61) function, which is not JAX compatible:

- The unitary matrix `U` is being converted to a NumPy array and copied
- The copied unitary matrix is being updated _in place_.
- NumPy functions (rather than `qml.math` functions) are used
- Exceptions based on value are being raised.


Pull request description:
Make `BasisRotation` jit compatible
**Context:**

PR #6019 only fixes `BasisRotation` when using backprop on `default.qubit`.  It is not jit compatible on any other device.  This is because `unitary_matrix` was being considered a hyperparameter, not a piece of data.  So we could not detect that the matrix was a tracer and we were in jitting mode, and we could not convert the matrix back into numpy data.

**Description of the Change:**

Make `unitary_matrix` a piece of data instead of a hyperparameter. This allows us to detect when it is being jitted.

As a by-product, I also made it valid pytree.

By making `unitary_matrix` a piece of data, we were able to get rid of the custom comparison method in `qml.equal`.

**Benefits:**

**Possible Drawbacks:**

**Related GitHub Issues:**

[sc-51603] Fixes #6004 

Patch:
diff --git a/doc/releases/changelog-0.40.0.md b/doc/releases/changelog-0.40.0.md
index e37693266..c1099ecf3 100644
--- a/doc/releases/changelog-0.40.0.md
+++ b/doc/releases/changelog-0.40.0.md
@@ -389,6 +389,7 @@
 
 * `qml.BasisRotation` template is now JIT compatible.
   [(#6019)](https://github.com/PennyLaneAI/pennylane/pull/6019)
+  [(#6779)](https://github.com/PennyLaneAI/pennylane/pull/6779)
 
 * The Jaxpr primitives for `for_loop`, `while_loop` and `cond` now store slices instead of
   numbers of args.
diff --git a/pennylane/ops/functions/equal.py b/pennylane/ops/functions/equal.py
index c7afc235b..f40769a7d 100644
--- a/pennylane/ops/functions/equal.py
+++ b/pennylane/ops/functions/equal.py
@@ -754,39 +754,6 @@ def _equal_counts(op1: CountsMP, op2: CountsMP, **kwargs):
     return _equal_measurements(op1, op2, **kwargs) and op1.all_outcomes == op2.all_outcomes
 
 
-@_equal_dispatch.register
-# pylint: disable=unused-argument
-def _equal_basis_rotation(
-    op1: qml.BasisRotation,
-    op2: qml.BasisRotation,
-    check_interface=True,
-    check_trainability=True,
-    rtol=1e-5,
-    atol=1e-9,
-):
-    if not qml.math.allclose(
-        op1.hyperparameters["unitary_matrix"],
-        op2.hyperparameters["unitary_matrix"],
-        atol=atol,
-        rtol=rtol,
-    ):
-        return (
-            "The hyperparameter unitary_matrix is not equal for op1 and op2.\n"
-            f"Got {op1.hyperparameters['unitary_matrix']}\n and {op2.hyperparameters['unitary_matrix']}."
-        )
-    if op1.wires != op2.wires:
-        return f"op1 and op2 have different wires. Got {op1.wires} and {op2.wires}."
-    if check_interface:
-        interface1 = qml.math.get_interface(op1.hyperparameters["unitary_matrix"])
-        interface2 = qml.math.get_interface(op2.hyperparameters["unitary_matrix"])
-        if interface1 != interface2:
-            return (
-                "The hyperparameter unitary_matrix has different interfaces for op1 and op2."
-                f" Got {interface1} and {interface2}."
-            )
-    return True
-
-
 @_equal_dispatch.register
 def _equal_hilbert_schmidt(
     op1: qml.HilbertSchmidt,
diff --git a/pennylane/templates/subroutines/basis_rotation.py b/pennylane/templates/subroutines/basis_rotation.py
index 832a06847..f522dcb8b 100644
--- a/pennylane/templates/subroutines/basis_rotation.py
+++ b/pennylane/templates/subroutines/basis_rotation.py
@@ -107,6 +107,10 @@ class BasisRotation(Operation):
 
         return cls._primitive.bind(*wires, unitary_matrix, check=check, id=id)
 
+    @classmethod
+    def _unflatten(cls, data, metadata):
+        return cls(wires=metadata[0], unitary_matrix=data[0])
+
     def __init__(self, wires, unitary_matrix, check=False, id=None):
         M, N = qml.math.shape(unitary_matrix)
 
@@ -124,19 +128,15 @@ class BasisRotation(Operation):
         if len(wires) < 2:
             raise ValueError(f"This template requires at least two wires, got {len(wires)}")
 
-        self._hyperparameters = {
-            "unitary_matrix": unitary_matrix,
-        }
-
-        super().__init__(wires=wires, id=id)
+        super().__init__(unitary_matrix, wires=wires, id=id)
 
     @property
     def num_params(self):
-        return 0
+        return 1
 
     @staticmethod
     def compute_decomposition(
-        wires, unitary_matrix, check=False
+        unitary_matrix, wires, check=False
     ):  # pylint: disable=arguments-differ
         r"""Representation of the operator as a product of other operators.
 


Test patch:
diff --git a/tests/ops/functions/test_equal.py b/tests/ops/functions/test_equal.py
index 2a6f52706..3ef0e9be2 100644
--- a/tests/ops/functions/test_equal.py
+++ b/tests/ops/functions/test_equal.py
@@ -2598,7 +2598,7 @@ class TestBasisRotation:
         assert_equal(op, other_op, atol=1e-5)
         assert qml.equal(op, other_op, rtol=0, atol=1e-9) is False
 
-        with pytest.raises(AssertionError, match="The hyperparameter unitary_matrix is not equal"):
+        with pytest.raises(AssertionError, match="op1 and op2 have different data"):
             assert_equal(op, other_op, rtol=0, atol=1e-9)
 
     @pytest.mark.parametrize("op, other_op", [(op1, op2)])
@@ -2629,7 +2629,7 @@ class TestBasisRotation:
         assert_equal(op, other_op, check_interface=False)
         assert qml.equal(op, other_op) is False
 
-        with pytest.raises(AssertionError, match=r"has different interfaces for op1 and op2"):
+        with pytest.raises(AssertionError, match=r"have different interfaces"):
             assert_equal(op, other_op)
 
 
diff --git a/tests/templates/test_subroutines/test_basis_rotation.py b/tests/templates/test_subroutines/test_basis_rotation.py
index a4328ca52..0f9f6412a 100644
--- a/tests/templates/test_subroutines/test_basis_rotation.py
+++ b/tests/templates/test_subroutines/test_basis_rotation.py
@@ -21,7 +21,6 @@ import pytest
 import pennylane as qml
 
 
-@pytest.mark.xfail  # to be fixed by [sc-51603]
 def test_standard_validity():
     """Run standard tests of operation validity."""
     weights = np.array(
@@ -402,8 +401,9 @@ class TestInterfaces:
 
         assert np.allclose(grads, np.zeros_like(unitary_matrix, dtype=complex), atol=tol, rtol=0)
 
+    @pytest.mark.parametrize("device_name", ("default.qubit", "reference.qubit"))
     @pytest.mark.jax
-    def test_jax_jit(self, tol):
+    def test_jax_jit(self, device_name, tol):
         """Test the jax interface."""
 
         import jax
@@ -417,7 +417,7 @@ class TestInterfaces:
             ]
         )
 
-        dev = qml.device("default.qubit", wires=3)
+        dev = qml.device(device_name, wires=3)
 
         circuit = jax.jit(qml.QNode(circuit_template, dev), static_argnames="check")
         circuit2 = qml.QNode(circuit_template, dev)