[Frontend] Implement an if-elif-else result type unification

[sergei-mironov]

Return values of conditional functions no longer need to be of exactly the same type. Catalyst would apply type promotion to branch return values if their types don't match.

[sc-41332]

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Comparison is base (5d818b7) 99.62% compared to head (6215cc7) 99.62%.

Additional details and impacted files

@@           Coverage Diff           @@
##             main     #333   +/-   ##
=======================================
  Coverage   99.62%   99.62%           
=======================================
  Files          42       42           
  Lines        7454     7480   +26     
  Branches      439      449   +10     
=======================================
+ Hits         7426     7452   +26     
  Misses         14       14           
  Partials       14       14           

Files	Coverage Δ
frontend/catalyst/jax_tracer.py	99.60% <100.00%> (+0.04%)	⬆️
frontend/catalyst/pennylane_extensions.py	99.64% <100.00%> (-0.01%)	⬇️
frontend/catalyst/utils/jax_extras.py	100.00% <100.00%> (ø)

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

[dime10]

Very nice, great work @grwlf 💯

[dime10]

Optional: Should we assert this?

[sergei-mironov]

We call the assertions in the _promote_jaxpr_types

[dime10]

Don't these assertions:

assert len(types) > 0, "Expected one or more set of types"
assert all(len(t) == len(types[0]) for t in types), "Expected matching number of arguments"

where types = [j.out_avals for j in jaxprs],

only check that the number of results from each jaxpr is the same? That is it doesn't check whether the pytrees match, nor whether array shapes match, which the docstring states as a requirement.

[sergei-mironov]

Ah, correct! I will add those.

[sergei-mironov]

I have added the shape assertion, but the pytree one is not that easy. We don't need the Pytrees when we do quantum tracing so there is no access to them. Adding them just to be able to assert here doesn't look feasible. We do check it elsewhere though.

[sergei-mironov]

But that means we don't need to mention pytrees in the comment so I have updated it.

[dime10]

I don't think it's important that the Pytrees match, otherwise we wouldn't know which Pytree to restore as the result of the conditional operation.

This would make for a good test case actually, where they match and where they don't.

[dime10]

Oh sorry, I misspelled. I meant "I do think it's important that the PyTrees match". I think we should add those test cases @grwlf

For example:

            @cond(True)
            def cond_fn():
                return (1, 1)

            @cond_fn.otherwise
            def cond_fn():
                return [2, 2]

should fail, because we don't know which PyTree to restore as the result of the operation, even though the number of arguments and dtypes match.

[sergei-mironov]

Addressed here #351