Fix the issue with dispatching to jax.vmap with qjit(vmap)

[maliasadi]

This PR fixes the issue with compiling catalyst.qjit(catalyst.vmap(fn))(*args, **kwargs)

For example, this example should pass now

dev = qml.device("lightning.qubit", wires=1)

@qml.qnode(dev)
def circuit(x, y):
    qml.RX(jnp.pi * x[0] + y, wires=0)
    qml.RY(x[1] ** 2, wires=0)
    qml.RX(x[1] * x[2], wires=0)
    return qml.expval(qml.PauliZ(0))

def cost(x, y, z):
    return circuit(x, y) * z

vmapped_cost = qjit(vmap(cost, in_axes=(0, 0, None)))

x = jnp.array([[0.1, 0.2, 0.3],
               [0.4, 0.5, 0.6],
               [0.7, 0.8, 0.9]])
y = jnp.array([jnp.pi, jnp.pi / 2, jnp.pi / 4])

vmapped_cost(x, y, 1)

[dime10]

Neat 👍 Maybe add this PR to the existing changelog entry?

[maliasadi]

@erick-xanadu Feel free to merge it.. 🙌

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 99.54%. Comparing base (1c754a3) to head (419747c).

Additional details and impacted files

@@            Coverage Diff             @@
##           v0.5.0-rc     #569   +/-   ##
==========================================
  Coverage      99.54%   99.54%           
==========================================
  Files             51       51           
  Lines           8618     8618           
  Branches         606      606           
==========================================
  Hits            8579     8579           
  Misses            21       21           
  Partials          18       18           

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

[josh146]

Can I merge this? :)

[erick-xanadu]

Can I merge this? :)

I want this one first: #568

After this, yes. I'll merge both of them as soon as possible and ping you @josh146 :)