Fix type promotion on branch (#1977)

**Context:**

```python
@qml.qjit(autograph=True)
@qml.qnode(qml.device('lightning.qubit', wires=1))
def circuit(cond1, cond2):
    if cond1:
        res = False
    elif cond2:
        res = 0.5
    else:
        res = False

    return res

circuit(False, True)
```

Should return value with type float64 instead of bool 

**Description of the Change:**

Choose the type that match the promoted type.

For an example:

```
branch_avals:    [[ShapedArray(float64[], weak_type=True)], [ShapedArray(bool[])]]
promoted_dtypes: [dtype('float64')]
```

Expected to return matched type `((ShapedArray(float64[],
weak_type=True), ...),)`


**Benefits:**

**Possible Drawbacks:**

**Related GitHub Issues:**

Fixes #1958 

[sc-97050]

---------

Co-authored-by: Joey Carter <[email protected]>
Co-authored-by: Paul <[email protected]>
Co-authored-by: David Ittah <[email protected]>

Author: 4 people

doc/releases/changelog-dev.md

@@ -60,6 +60,9 @@
 
 <h3>Bug fixes 🐛</h3>
 
+* Fix type promotion on conditional branches, where the return values from `cond` should be the promoted one.
+  [(#1977)](https://github.com/PennyLaneAI/catalyst/pull/1977)
+
 * Fix wrong handling of partitioned shots in the decomposition pass of `measurements_from_samples`.
   [(#1981)](https://github.com/PennyLaneAI/catalyst/pull/1981)
 

frontend/catalyst/api_extensions/control_flow.py

@@ -60,6 +60,7 @@
     HybridOp,
     HybridOpRegion,
     QRegPromise,
+    _promote_jaxpr_types,
     trace_function,
     trace_quantum_operations,
     unify_convert_result_types,
@@ -162,7 +163,8 @@ def cond_else():
 
     The conditional function is permitted to also return values.
     Any value that is supported by JAX JIT compilation is supported as a return
-    type.
+    type. If provided, return types need to be identical or at least promotable across both
+    branches.
 
     .. code-block:: python
 
@@ -185,31 +187,22 @@ def conditional_fn():
         There are various constraints and restrictions that should be kept in mind
         when working with conditionals in Catalyst.
 
-        The return values of all branches of :func:`~.cond` must be the same type.
-        Returning different types, or ommitting a return value in one branch (e.g.,
+        The return values of all branches of :func:`~.cond` must be the same shape.
+        Returning different shapes, or ommitting a return value in one branch (e.g.,
         returning ``None``) but not in others will result in an error.
 
+        However, the return values of all branches of :func:`~.cond` can be different data types.
+        In this case, the return types will automatically be promoted to the next common larger
+        type.
+
         >>> @qjit
         ... def f(x: float):
         ...     @cond(x > 1.5)
         ...     def cond_fn():
         ...         return x ** 2  # float
         ...     @cond_fn.otherwise
         ...     def else_branch():
-        ...         return 6  # int
-        ...     return cond_fn()
-        TypeError: Conditional requires consistent return types across all branches, got:
-        - Branch at index 0: [ShapedArray(float64[], weak_type=True)]
-        - Branch at index 1: [ShapedArray(int64[], weak_type=True)]
-        Please specify an else branch if none was specified.
-        >>> @qjit
-        ... def f(x: float):
-        ...     @cond(x > 1.5)
-        ...     def cond_fn():
-        ...         return x ** 2  # float
-        ...     @cond_fn.otherwise
-        ...     def else_branch():
-        ...         return 6.  # float
+        ...         return 6  # int (promotable to float)
         ...     return cond_fn()
         >>> f(1.5)
         Array(6., dtype=float64)
@@ -224,10 +217,9 @@ def conditional_fn():
         ...     def cond_fn():
         ...         return x ** 2
         ...     return cond_fn()
-        TypeError: Conditional requires consistent return types across all branches, got:
-        - Branch at index 0: [ShapedArray(float64[], weak_type=True)]
-        - Branch at index 1: []
-        Please specify an else branch if none was specified.
+        TypeError: Conditional requires a consistent return structure across all branches! Got
+        PyTreeDef(None) and PyTreeDef(*). Please specify an else branch if PyTreeDef(None) was
+        specified.
 
         >>> @qjit
         ... def f(x: float):
@@ -774,18 +766,16 @@ def _call_with_quantum_ctx(self):
         out_tree = out_sigs[-1].out_tree()
         all_consts = [s.out_consts() for s in out_sigs]
         out_types = [s.out_type() for s in out_sigs]
-        # FIXME: We want to perform the result unificaiton here:
-        # all_jaxprs = [s.out_initial_jaxpr() for s in out_sigs]
-        # all_noimplouts = [s.num_implicit_outputs() for s in out_sigs]
-        # _, out_type, _, all_consts = unify_convert_result_types(
-        #     all_jaxprs, all_consts, all_noimplouts
-        # )
-        # Unfortunately, we can not do this beacuse some tracers (specifically, the results of
-        # ``qml.measure``) might not have their source Jaxpr equation yet. Thus, we delay the
-        # unification until the quantum tracing is done. The consequence of that: we have to guess
-        # the output type now and if we fail to do so, we might face MLIR type error down the
-        # pipeline.
+
+        # Select the output type of the one with the promoted dtype among all branches
         out_type = out_types[-1]
+        branch_avals = [[aval for aval, _ in branch_out_type] for branch_out_type in out_types]
+        promoted_dtypes = _promote_jaxpr_types(branch_avals)
+
+        out_type = [
+            (aval.update(dtype=dtype), expl)
+            for dtype, (aval, expl) in zip(promoted_dtypes, out_type)
+        ]
 
         # Create output tracers in the outer tracing context
         out_expanded_classical_tracers = output_type_to_tracers(
@@ -1563,14 +1553,20 @@ def _assert_cond_result_structure(trees: List[PyTreeDef]):
         if tree != expected_tree:
             raise TypeError(
                 "Conditional requires a consistent return structure across all branches! "
-                f"Got {tree} and {expected_tree}."
+                f"Got {tree} and {expected_tree}. Please specify an else branch if PyTreeDef(None) "
+                "was specified."
             )
 
 
 def _assert_cond_result_types(signatures: List[List[AbstractValue]]):
     """Ensure a consistent type signature across branch results."""
     num_results = len(signatures[0])
-    assert all(len(sig) == num_results for sig in signatures), "mismatch: number or results"
+
+    if not all(len(sig) == num_results for sig in signatures):
+        raise TypeError(
+            "Conditional requires a consistent number of results across all branches! It might "
+            "happen when some of the branch returns dynamic shape and some return constant shape."
+        )
 
     for i in range(num_results):
         aval_slice = [avals[i] for avals in signatures]

frontend/test/pytest/test_conditionals.py

@@ -14,6 +14,7 @@
 
 from textwrap import dedent
 
+import jax
 import jax.numpy as jnp
 import numpy as np
 import pennylane as qml
@@ -64,6 +65,7 @@ def asline(text):
         assert asline(expected) == asline(circuit.jaxpr)
 
 
+# pylint: disable=too-many-public-methods,too-many-lines
 class TestCond:
     """Test suite for the Cond functionality in Catalyst."""
 
@@ -311,6 +313,108 @@ def cond_else():
 
         assert 0 == circuit()
 
+    def test_branch_multi_return_type_unification_qjit_2(self):
+        """Test that unification happens before the results of the cond primitve is available."""
+
+        @qjit
+        def circuit(cond1, cond2):
+            @cond(cond1)
+            def cond_fn():
+                return False
+
+            @cond_fn.else_if(cond2)
+            def cond_fn_2():
+                return 0.5
+
+            @cond_fn.otherwise
+            def cond_fn_3():
+                return False
+
+            r = cond_fn()
+            assert r.dtype is jnp.dtype(
+                "float64" if jax.config.values["jax_enable_x64"] else "float32"
+            )
+            return r
+
+        assert 0.5 == circuit(False, True)
+
+    def test_branch_multi_return_type_unification_qjit_3(self):
+        """Test that unification happens before the results of the cond primitve is available."""
+
+        @qjit
+        def circuit(cond1, cond2):
+            @cond(cond1)
+            def cond_fn():
+                return False
+
+            @cond_fn.else_if(cond2)
+            def cond_fn_2():
+                return False
+
+            @cond_fn.otherwise
+            def cond_fn_3():
+                return 0.5
+
+            r = cond_fn()
+            assert r.dtype is jnp.dtype(
+                "float64" if jax.config.values["jax_enable_x64"] else "float32"
+            )
+            return r
+
+        assert 0.0 == circuit(False, True)
+
+    def test_branch_multi_return_type_unification_qjit_4(self):
+        """Test that unification happens before the results of the cond primitve is available."""
+
+        @qjit
+        def circuit(cond1, cond2):
+            @cond(cond1)
+            def cond_fn():
+                return {0: True, 1: 0.5}
+
+            @cond_fn.else_if(cond2)
+            def cond_fn_2():
+                return {0: 0.7, 1: True}
+
+            @cond_fn.otherwise
+            def cond_fn_3():
+                return {0: True, 1: False}
+
+            r = cond_fn()
+            expected_dtype = jnp.dtype(
+                "float64" if jax.config.values["jax_enable_x64"] else "float32"
+            )
+            assert all(v.dtype is expected_dtype for _, v in r.items())
+            return r
+
+        assert {0: 0.7, 1: 1.0} == circuit(False, True)
+
+    def test_qnode_cond_inconsistent_return_types(self, backend):
+        """Test that catalyst raises an error when the conditional has inconsistent return types."""
+
+        @qjit
+        @qml.qnode(qml.device(backend, wires=4))
+        def f(flag, sz):
+            a = jnp.ones([sz], dtype=float)
+            b = jnp.zeros([3], dtype=float)
+
+            @cond(flag)
+            def case():
+                return a
+
+            @case.otherwise
+            def case():
+                return b
+
+            c = case()
+            return c
+
+        with pytest.raises(
+            TypeError,
+            match="Conditional requires a consistent number of results across all branches",
+        ):
+            f(True, 3)
+
     @pytest.mark.xfail(
         reason="Inability to apply Jax transformations before the quantum traing is complete"
     )