Support pytree inputs when program capture enabled (#2165)

**Context:**

Change broken off from #2078, since it's basically it's own feature.

**Description of the Change:**

Just determine the inputs to `from_plxpr` from the plxpr itself. This
means they will always match what's expected.

**Benefits:**

Always accurate with arbitrary nested arguments.

**Possible Drawbacks:**

Won't work with dynamically shaped inputs, but we don't support that yet
anyway.

**Related GitHub Issues:**

[sc-102817]

Author: albi3ro

doc/releases/changelog-dev.md

@@ -12,6 +12,9 @@
   :func:`~.qjit` with program capture enabled.
   [(#2154)](https://github.com/PennyLaneAI/catalyst/pull/2154)
 
+* Pytree inputs can now be used when program capture is enabled.
+  [(#2165)](https://github.com/PennyLaneAI/catalyst/pull/2165)
+
 <h3>Breaking changes 💔</h3>
 
 <h3>Deprecations 👋</h3>

frontend/catalyst/from_plxpr/from_plxpr.py

@@ -399,7 +399,7 @@ def trace_from_pennylane(
             fn.static_argnums = static_argnums
 
         plxpr, out_type, out_treedef = make_jaxpr2(fn, **make_jaxpr_kwargs)(*args, **kwargs)
-        jaxpr = from_plxpr(plxpr)(*dynamic_args, **kwargs)
+        jaxpr = from_plxpr(plxpr)(*plxpr.in_avals)
 
     return jaxpr, out_type, out_treedef, sig
 

frontend/test/pytest/test_pytree_args.py

@@ -22,13 +22,15 @@
 import pennylane as qml
 import pytest
 from jax._src.tree_util import tree_flatten
+from pennylane import adjoint, cond, for_loop, grad, qjit
 
-from catalyst import adjoint, cond, for_loop, grad, measure, qjit
+from catalyst import measure
 
 
 class TestPyTreesReturnValues:
     """Test QJIT workflows with different return value data-types."""
 
+    @pytest.mark.usefixtures("use_both_frontend")
     def test_return_value_float(self, backend):
         """Test constant."""
 
@@ -41,10 +43,13 @@ def circuit1(params):
         jitted_fn = qjit(circuit1)
 
         params = [0.4, 0.8]
-        expected = 0.64170937
+        expected = jnp.cos(params[0]) * jnp.cos(params[1])
         result = jitted_fn(params)
         assert jnp.allclose(result, expected)
 
+    def test_return_value_mcm(self, backend):
+        """Test that a qnode can return a scalar mcm."""
+
         @qml.qnode(qml.device(backend, wires=2))
         def circuit2():
             return measure(0)
@@ -166,6 +171,7 @@ def _f(x):
         assert result[0] == 4.0
         assert result[1] == 6.0
 
+    @pytest.mark.usefixtures("use_both_frontend")
     def test_return_value_hybrid(self, backend):
         """Test tuples."""
 
@@ -239,6 +245,9 @@ def else_fn():
     def test_return_value_dict(self, backend, tol_stochastic, mcm_method):
         """Test dictionaries."""
 
+        if mcm_method == "one-shot" and qml.capture.enabled():
+            pytest.xfail()
+
         @qml.qnode(qml.device(backend, wires=2))
         def circuit1(params):
             qml.RX(params[0], wires=0)
@@ -438,6 +447,7 @@ def circuit1(params):
         }
         result = jitted_fn(params)
 
+    @pytest.mark.usefixtures("use_both_frontend")
     def test_args_workflow(self, backend):
         """Test arguments with workflows."""
 
@@ -527,6 +537,7 @@ def workflow2(params):
         assert np.allclose(result_flatten, result_flatten_expected)
         assert tree == tree_expected
 
+    @pytest.mark.usefixtures("use_both_frontend")
     @pytest.mark.parametrize("inp", [(np.array([0.2, 0.5])), (jnp.array([0.2, 0.5]))])
     def test_args_control_flow(self, backend, inp):
         """Test arguments with control-flows operations."""
@@ -574,6 +585,7 @@ def circuit(dictionary):
         result = circuit({"wire": 1})
         assert jnp.allclose(result, True)
 
+    @pytest.mark.usefixtures("use_both_frontend")
     def test_dev_wires_have_pytree(self, backend):
         """Device wires are pytree-compatible."""
 
@@ -592,6 +604,7 @@ def test_function():
         test_function()
 
 
+@pytest.mark.usefixtures("use_both_frontend")
 class TestAuxiliaryData:
     """Test PyTrees with Auxiliary data."""