Add support for JIT when using DiagonalQubitUnitary

[josh146]

Context: Due to a unitarity check, qml.DiagonalQubitUnitary did not work with @jax.jit and @tf.function.

Description of the Change: Uses qml.math.is_abstract to avoid the unitarity check if a JIT is detected.

Benefits: The operation now works with JIT compilation.

Possible Drawbacks: TensorFlow is very slow to compile, so the test added to ensure that @tf.function works takes 12 seconds :( Should we remove it?

Related GitHub Issues: n/a

[codecov]

Codecov Report

Merging #2445 (dc6a107) into master (a64c204) will not change coverage.
The diff coverage is 100.00%.

@@           Coverage Diff           @@
##           master    #2445   +/-   ##
=======================================
  Coverage   99.46%   99.46%           
=======================================
  Files         242      242           
  Lines       19104    19104           
=======================================
  Hits        19002    19002           
  Misses        102      102           

Impacted Files	Coverage Δ
pennylane/ops/qubit/matrix_ops.py	99.12% <100.00%> (ø)

---

Continue to review full report at Codecov.

Legend - Click here to learn more
Δ = absolute <relative> (impact), ø = not affected, ? = missing data
Powered by Codecov. Last update a64c204...dc6a107. Read the comment docs.

[antalszava]

Looks good! 🎉 😍