[Program Capture] higher order primitives support pytree input and output

[albi3ro]

Context:

Similar to jax, we should be able to catch higher order primitives that accept arbitrary pytree inputs and outputs. Currently, our primitives assume flat inputs and flat outputs.

Description of the Change:

Updates qnode_call, for_loop, and cond to allow arbitrary pytree inputs and outputs.

For example,

@qml.for_loop(1, 7, 2)
def f(i, x):
    return {"x": i+x["x"]}

x = {"x": 0}
f(x) 

{'x': Array(9, dtype=int32, weak_type=True)}

Benefits:

Easier to cart complicated combinations of variables around. More versatile program capture.

Possible Drawbacks:

If different branches (ie cond) have different pytree structures, we may run into weird behavior.

Related GitHub Issues:

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 99.65%. Comparing base (5a7a163) to head (dd8f722).
Report is 1 commits behind head on master.

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #6081      +/-   ##
==========================================
- Coverage   99.66%   99.65%   -0.01%     
==========================================
  Files         430      431       +1     
  Lines       41857    41593     -264     
==========================================
- Hits        41716    41451     -265     
- Misses        141      142       +1     

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

[mudit2812]

Looks good. Now that the while loop capture is merged, maybe worth adding pytree capture support for that as well?

[mudit2812]

I'm not going to block approval, but could you also add a test for qml.cond with MCM predicates that accepts pytree arguments? Conditionals with MCMs cannot return anything, so the only thing that would need to be tested is pytree inputs.

[obliviateandsurrender]

Looks good, just one small non-blocking doubt!

[albi3ro]

I'm not going to block approval, but could you also add a test for qml.cond with MCM predicates that accepts pytree arguments? Conditionals with MCMs cannot return anything, so the only thing that would need to be tested is pytree inputs.

Added.