Repository: PennyLaneAI/pennylane

Issue:
[BUG] Error when passing complex gate arg with JAX, backprop and the beta QNode
### Expected behavior

The following executes well:
```python
import pennylane as qml
from pennylane import numpy as np
from jax import numpy as jnp
import jax

dev = qml.device('default.qubit', wires=2)

@qml.beta.qnode(dev, diff_method="backprop", interface='jax')
def test(x):
    qml.RZ(x, wires=[0])
    return qml.state()

jax.jacobian(test,holomorphic=True)(0.3 + 0j)
```

### Actual behavior

An error is raised:
```
~/anaconda3/lib/python3.8/site-packages/jax/interpreters/xla.py in backend_compile(backend, built_c, options)
    367   # we use a separate function call to ensure that XLA compilation appears
    368   # separately in Python profiling results
--> 369   return backend.compile(built_c, compile_options=options)
    370 
    371 

RuntimeError: UNIMPLEMENTED: binary complex op 'remainder'
```

### Additional information

JAX seems to not provide the gradient for `%`.

### Source code

_No response_

### Tracebacks

_No response_

### System information

```shell
Name: PennyLane
Version: 0.19.0.dev0
Summary: PennyLane is a Python quantum machine learning library by Xanadu Inc.
Home-page: https://github.com/XanaduAI/pennylane
Author: None
Author-email: None
License: Apache License 2.0
Location: /home/antal/xanadu/pennylane
Requires: numpy, scipy, networkx, autograd, toml, appdirs, semantic-version, autoray, cachetools, pennylane-lightning
Required-by: PennyLane-Cirq, PennyLane-Orquestra, PennyLane-SF, pennylane-qulacs, PennyLane-IonQ, amazon-braket-pennylane-plugin, PennyLane-Forest, PennyLane-Honeywell, PennyLane-qiskit, PennyLane-AQT, PennyLane-Qchem, PennyLane-PQ, PennyLane-Lightning
Platform info:           Linux-5.11.0-37-generic-x86_64-with-glibc2.10
Python version:          3.8.5
Numpy version:           1.20.3
Scipy version:           1.7.1
Installed devices:
- cirq.mixedsimulator (PennyLane-Cirq-0.17.1)
- cirq.pasqal (PennyLane-Cirq-0.17.1)
- cirq.qsim (PennyLane-Cirq-0.17.1)
- cirq.qsimh (PennyLane-Cirq-0.17.1)
- cirq.simulator (PennyLane-Cirq-0.17.1)
- orquestra.forest (PennyLane-Orquestra-0.15.0)
- orquestra.ibmq (PennyLane-Orquestra-0.15.0)
- orquestra.qiskit (PennyLane-Orquestra-0.15.0)
- orquestra.qulacs (PennyLane-Orquestra-0.15.0)
- strawberryfields.fock (PennyLane-SF-0.16.0.dev0)
- strawberryfields.gaussian (PennyLane-SF-0.16.0.dev0)
- strawberryfields.gbs (PennyLane-SF-0.16.0.dev0)
- strawberryfields.remote (PennyLane-SF-0.16.0.dev0)
- strawberryfields.tf (PennyLane-SF-0.16.0.dev0)
- qulacs.simulator (pennylane-qulacs-0.17.0.dev0)
- ionq.qpu (PennyLane-IonQ-0.17.0.dev0)
- ionq.simulator (PennyLane-IonQ-0.17.0.dev0)
- braket.aws.qubit (amazon-braket-pennylane-plugin-1.4.1.dev0)
- braket.local.qubit (amazon-braket-pennylane-plugin-1.4.1.dev0)
- forest.numpy_wavefunction (PennyLane-Forest-0.17.0.dev0)
- forest.qvm (PennyLane-Forest-0.17.0.dev0)
- forest.wavefunction (PennyLane-Forest-0.17.0.dev0)
- honeywell.hqs (PennyLane-Honeywell-0.16.0.dev0)
- qiskit.aer (PennyLane-qiskit-0.18.0.dev0)
- qiskit.basicaer (PennyLane-qiskit-0.18.0.dev0)
- qiskit.ibmq (PennyLane-qiskit-0.18.0.dev0)
- aqt.noisy_sim (PennyLane-AQT-0.18.0)
- aqt.sim (PennyLane-AQT-0.18.0)
- projectq.classical (PennyLane-PQ-0.18.0.dev0)
- projectq.ibm (PennyLane-PQ-0.18.0.dev0)
- projectq.simulator (PennyLane-PQ-0.18.0.dev0)
- lightning.qubit (PennyLane-Lightning-0.19.0.dev0)
- default.gaussian (PennyLane-0.19.0.dev0)
- default.mixed (PennyLane-0.19.0.dev0)
- default.qubit (PennyLane-0.19.0.dev0)
- default.qubit.autograd (PennyLane-0.19.0.dev0)
- default.qubit.jax (PennyLane-0.19.0.dev0)
- default.qubit.tf (PennyLane-0.19.0.dev0)
- default.qubit.torch (PennyLane-0.19.0.dev0)
- default.tensor (PennyLane-0.19.0.dev0)
- default.tensor.tf (PennyLane-0.19.0.dev0)
```
```


### 

- [X] I have searched exisisting GitHub issues to make sure the issue does not already exist.


Discussion:

Resolving patch:
diff --git a/pennylane/operation.py b/pennylane/operation.py
--- a/pennylane/operation.py
+++ b/pennylane/operation.py
@@ -235,11 +235,15 @@ def classproperty(func):
 
 
 def _process_data(op):
+
+    # Use qml.math.real to take the real part. We may get complex inputs for
+    # example when differentiating holomorphic functions with JAX: a complex
+    # valued QNode (one that returns qml.state) requires complex typed inputs.
     if op.name in ("RX", "RY", "RZ", "PhaseShift", "Rot"):
-        return str([qml.math.round(d % (2 * np.pi), 10) for d in op.data])
+        return str([qml.math.round(qml.math.real(d) % (2 * np.pi), 10) for d in op.data])
 
     if op.name in ("CRX", "CRY", "CRZ", "CRot"):
-        return str([qml.math.round(d % (4 * np.pi), 10) for d in op.data])
+        return str([qml.math.round(qml.math.real(d) % (4 * np.pi), 10) for d in op.data])
 
     return str(op.data)
 


Test patch:
diff --git a/tests/ops/qubit/test_parametric_ops.py b/tests/ops/qubit/test_parametric_ops.py
--- a/tests/ops/qubit/test_parametric_ops.py
+++ b/tests/ops/qubit/test_parametric_ops.py
@@ -873,6 +873,23 @@ def circuit(phi):
         res = tape.gradient(result, phi)
         assert np.allclose(res, expected, atol=tol, rtol=0)
 
+    @pytest.mark.parametrize("par", np.linspace(0, 2 * np.pi, 3))
+    def test_qnode_with_rx_and_state_jacobian_jax(self, par, tol):
+        """Test the jacobian of a complex valued QNode that contains a rotation
+        using the JAX interface."""
+        jax = pytest.importorskip("jax")
+
+        dev = qml.device("default.qubit", wires=1)
+
+        @qml.qnode(dev, diff_method="backprop", interface="jax")
+        def test(x):
+            qml.RX(x, wires=[0])
+            return qml.state()
+
+        res = jax.jacobian(test, holomorphic=True)(par + 0j)
+        expected = -1 / 2 * np.sin(par / 2), -1 / 2 * 1j * np.cos(par / 2)
+        assert np.allclose(res, expected, atol=tol, rtol=0)
+
 
 PAULI_ROT_PARAMETRIC_MATRIX_TEST_DATA = [
     (