Make qml.QutritBasisStatePreparation JIT compatible

[willjmax]

This PR makes qml.QutritBasisStatePreparation JIT compatible. Previously, the template used the qml.TShift operator in the decomposition. However, this approach requires non-jittable control flow on the input. To make the template jittable the decomposition is changed to return a qml.ops.QutritUnitary representing a TShift operator raised to some power. This change requires adding matrix_power and eigh to the multi dispatch. This PR addresses sc-70863.

[github-actions]

Hello. You may have forgotten to update the changelog!
Please edit doc/releases/changelog-dev.md with:

• A one-to-two sentence description of the change. You may include a small working example for new features.
• A link back to this PR.
• Your name (or GitHub username) in the contributors section.

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 99.39%. Comparing base (38a5140) to head (111405e).
Report is 321 commits behind head on master.

Additional details and impacted files

@@           Coverage Diff           @@
##           master    #6308   +/-   ##
=======================================
  Coverage   99.39%   99.39%           
=======================================
  Files         448      448           
  Lines       42464    42471    +7     
=======================================
+ Hits        42209    42216    +7     
  Misses        255      255           

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[obliviateandsurrender]

Thanks @willjmax! Some minor queries, before I dig in deeper.

[KetpuntoG]

Nice! 😄

Just one comment. The following example does not work with jax.jit. For some reason qml.sample fails. I have tested default.qubit with BasisState to see if it should fail but it works correctly there.

n = 2

@jax.jit
@qml.qnode(qml.device("default.qutrit", wires = n, shots = 1))
def circuit(state):

    qml.QutritBasisStatePreparation(state, wires = range(n))
    
    return qml.sample(wires = range(n))

state = jax.numpy.array([1,1])
circuit(state)

Otherwise, the PR looks good to me :)

[KetpuntoG]

Looks good! 👌

[obliviateandsurrender]

[Non blocker] could combine matrix and pl_gates tests into one decomposition test, otherwise looks good with just one minor thing.

[obliviateandsurrender]

I think qutrit basis state prep test in templates/state_preparations/basis_qutrit.py should be made to xfail for the CI to pass.

[willjmax]

I think qutrit basis state prep test in templates/state_preparations/basis_qutrit.py should be made to xfail for the CI to pass.

No, this was due to a bug that I caused. I pushed a fix.