[Capture] add jit and vmap support to eval_jaxpr

[albi3ro]

Note that this change isn't necessarily high priority, but I thought I'd just open the PR since I already had the code.

Context:

The lightning devices currently throw an error with jax.jit and jax.vmap and program capture. While there might be some discussion about where the jax.pure_callback should like long term, for now I'm just adding the jax.pure_callback call inside the relevant devices. Not many devices will be implementing Device.eval_jaxpr at the moment anyway.

Description of the Change:

Adds a jax.pure_callback layer in Device.eval_jaxpr.

Benefits:

jax.jit and jax.vmap can now be used with capture and lightning devices.

Possible Drawbacks:

I don't know if we want to have the pure_callback layer live on the device long term.

Also, it turns all errors into jaxlib.xla_extension.XlaRuntimeError, which a bit scarier to debug.

Related Shortcut Stories:
[sc-83589]

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 66.30%. Comparing base (7d24015) to head (046155c).
Report is 1 commits behind head on master.

❗ There is a different number of reports uploaded between BASE (7d24015) and HEAD (046155c). Click for more details.

HEAD has 29 uploads less than BASE

Flag	BASE (7d24015)	HEAD (046155c)
	35	6

Additional details and impacted files

@@             Coverage Diff             @@
##           master    #1055       +/-   ##
===========================================
- Coverage   98.00%   66.30%   -31.71%     
===========================================
  Files         233       28      -205     
  Lines       39642     2511    -37131     
===========================================
- Hits        38853     1665    -37188     
- Misses        789      846       +57     

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

[mudit2812]

Ready to approve except for the testing comment.

[JerryChen97]

Looks good to me, though probably we also need to hear from Lightning team to see if any concerns regarding the dependecy introduced @josephleekl @maliasadi

[JerryChen97]

LGTM!

[AmintorDusko]

Are these codecov warnings all false positives?

[albi3ro]

Are these codecov warnings all false positives?

I also got them on #1067 as well. These tests are actually running, since I got failures in the CI I had to fix. So I don't really know what's causing code cov to be so confused.

[albi3ro]

UPDATE: FIXED IT: I was being stupid.

So I'm seeing in the CI:

FAILED tests/test_eval_jaxpr.py::test_vmap_in_axes[0-1] - assert False

but when I run:

in_axis=0
out_axis=1

@qml.qnode(qml.device('lightning.qubit', wires=1))
def circuit(mat):
    qml.QubitUnitary(mat, 0)
    return qml.expval(qml.Z(0)), qml.state()

mats = jax.numpy.stack(
    [qml.X.compute_matrix(), qml.Y.compute_matrix(), qml.Z.compute_matrix()], axis=in_axis
)
expval, state = jax.vmap(circuit, in_axes=in_axis, out_axes=(0, out_axis))(mats)

assert expval.shape == (3,)
assert qml.math.allclose(expval, jax.numpy.array([-1, -1, 1]))  # flip, flip, no flip

assert state.shape == (3, 4) if out_axis == 0 else (4, 3)

It works fine. If my install was bad, I would expect this to error out.

I'm confused.

[maliasadi]

Thanks @albi3ro! Happy to approve