Daily rc sync to main (#1647)

Automatic sync from the release candidate to main during a feature
freeze.

---------

Co-authored-by: Mehrdad Malek <[email protected]>
Co-authored-by: erick-xanadu <[email protected]>
Co-authored-by: Raul Torres <[email protected]>
Co-authored-by: Sengthai Heng <[email protected]>
Co-authored-by: David Ittah <[email protected]>
Co-authored-by: Paul <[email protected]>
Co-authored-by: Joey Carter <[email protected]>
Co-authored-by: GitHub Actions Bot <>
Co-authored-by: Mehrdad Malekmohammadi <[email protected]>

Author: 8 people

doc/releases/changelog-0.11.0.md

@@ -223,6 +223,7 @@
 
 * Added `qjit.mlir_opt` property to access the optimized MLIR representation of a compiled function.
   [(#1579)](https://github.com/PennyLaneAI/catalyst/pull/1579)
+  [(#1637)](https://github.com/PennyLaneAI/catalyst/pull/1637)
 
 * Improve error message for ZNE.
   [(#1603)](https://github.com/PennyLaneAI/catalyst/pull/1603)
@@ -268,6 +269,9 @@
   pytrees inside a loop with autograph causes falling back to python.
   [(#1601)](https://github.com/PennyLaneAI/catalyst/pull/1601)
 
+* Closure variables are now supported with `grad` and `value_and_grad`.
+  [(#1613)](https://github.com/PennyLaneAI/catalyst/pull/1613)
+
 <h3>Internal changes ⚙️</h3>
 
 * Updated the call signature for the PLXPR `qnode_prim` primitive.

frontend/catalyst/api_extensions/differentiation.py

@@ -667,7 +667,7 @@ def __call__(self, *args, **kwargs):
 
                     # It always returns list as required by catalyst control-flows
                     results = value_and_grad_p.bind(
-                        *input_data_flat, jaxpr=jaxpr, fn=fn, grad_params=grad_params
+                        *input_data_flat, *jaxpr.consts, jaxpr=jaxpr, fn=fn, grad_params=grad_params
                     )
 
                     # value_and_grad returns two results: the values and the gradients,
@@ -686,7 +686,7 @@ def __call__(self, *args, **kwargs):
 
                     # It always returns list as required by catalyst control-flows
                     results = grad_p.bind(
-                        *input_data_flat, jaxpr=jaxpr, fn=fn, grad_params=grad_params
+                        *input_data_flat, *jaxpr.consts, jaxpr=jaxpr, fn=fn, grad_params=grad_params
                     )
 
                     # grad returns only the gradients,

frontend/catalyst/compiler.py

@@ -290,8 +290,11 @@ def _catalyst(*args, stdin=None):
     catalyst *args
     """
     cmd = _get_catalyst_cli_cmd(*args, stdin=stdin)
-    result = subprocess.run(cmd, input=stdin, check=True, capture_output=True, text=True)
-    return result.stdout
+    try:
+        result = subprocess.run(cmd, input=stdin, check=True, capture_output=True, text=True)
+        return result.stdout
+    except subprocess.CalledProcessError as e:
+        raise CompileError(f"catalyst failed with error code {e.returncode}: {e.stderr}") from e
 
 
 def _quantum_opt(*args, stdin=None):

frontend/catalyst/jax_primitives.py

@@ -500,6 +500,27 @@ def _grad_lowering(ctx, *args, jaxpr, fn, grad_params):
         argnums: argument indices which define over which arguments to
             differentiate.
     """
+    consts = []
+    offset = len(args) - len(jaxpr.consts)
+    for i, jax_array_or_tracer in enumerate(jaxpr.consts):
+        if not isinstance(
+            jax_array_or_tracer, jax._src.interpreters.partial_eval.DynamicJaxprTracer
+        ):
+            # ``ir.DenseElementsAttr.get()`` constructs a dense elements attribute from an array of
+            # element values. This doesn't support ``jaxlib.xla_extension.Array``, so we have to
+            # cast such constants to numpy array types.
+            const = jax_array_or_tracer
+            const_type = shape_dtype_to_ir_type(const.shape, const.dtype)
+            nparray = np.asarray(const)
+            attr = ir.DenseElementsAttr.get(nparray, type=const_type)
+            constval = StableHLOConstantOp(attr).results
+            consts.append(constval)
+        else:
+            # There are some cases where this value cannot be converted into
+            # a jax.numpy.array.
+            # in that case we get it from the arguments.
+            consts.append(args[offset + i])
+
     method, h, argnums = grad_params.method, grad_params.h, grad_params.expanded_argnums
     mlir_ctx = ctx.module_context.context
     finiteDiffParam = None
@@ -516,18 +537,9 @@ def _grad_lowering(ctx, *args, jaxpr, fn, grad_params):
     output_types = list(map(mlir.aval_to_ir_types, ctx.avals_out))
     flat_output_types = util.flatten(output_types)
 
-    # ``ir.DenseElementsAttr.get()`` constructs a dense elements attribute from an array of
-    # element values. This doesn't support ``jaxlib.xla_extension.Array``, so we have to cast
-    # such constants to numpy array types.
-
-    constants = []
-    for const in jaxpr.consts:
-        const_type = shape_dtype_to_ir_type(const.shape, const.dtype)
-        nparray = np.asarray(const)
-        attr = ir.DenseElementsAttr.get(nparray, type=const_type)
-        constantVals = StableHLOConstantOp(attr).results
-        constants.append(constantVals)
-    args_and_consts = constants + list(args)
+    len_args = len(args)
+    index = len_args - len(consts)
+    args_and_consts = consts + list(args[:index])
 
     return GradOp(
         flat_output_types,
@@ -563,21 +575,38 @@ def _value_and_grad_lowering(ctx, *args, jaxpr, fn, grad_params):
     Returns:
         MLIR results
     """
-    args = list(args)
+    consts = []
+    offset = len(args) - len(jaxpr.consts)
+    for i, jax_array_or_tracer in enumerate(jaxpr.consts):
+        if not isinstance(
+            jax_array_or_tracer, jax._src.interpreters.partial_eval.DynamicJaxprTracer
+        ):
+            # ``ir.DenseElementsAttr.get()`` constructs a dense elements attribute from an array of
+            # element values. This doesn't support ``jaxlib.xla_extension.Array``, so we have to
+            # cast such constants to numpy array types.
+            const = jax_array_or_tracer
+            const_type = shape_dtype_to_ir_type(const.shape, const.dtype)
+            nparray = np.asarray(const)
+            attr = ir.DenseElementsAttr.get(nparray, type=const_type)
+            constval = StableHLOConstantOp(attr).results
+            consts.append(constval)
+        else:
+            # There are some cases where this value cannot be converted into
+            # a jax.numpy.array.
+            # in that case we get it from the arguments.
+            consts.append(args[offset + i])
+
+    len_args = len(args)
+    index = len_args - len(consts)
+    args = list(args[0:index])
     method, h, argnums = grad_params.method, grad_params.h, grad_params.expanded_argnums
     mlir_ctx = ctx.module_context.context
     new_argnums = np.array([len(jaxpr.consts) + num for num in argnums])
 
     output_types = list(map(mlir.aval_to_ir_types, ctx.avals_out))
     flat_output_types = util.flatten(output_types)
 
-    constants = []
-    for const in jaxpr.consts:
-        const_type = shape_dtype_to_ir_type(const.shape, const.dtype)
-        nparray = np.asarray(const)
-        attr = ir.DenseElementsAttr.get(nparray, type=const_type)
-        constantVals = StableHLOConstantOp(attr).results
-        constants.append(constantVals)
+    constants = consts
 
     consts_and_args = constants + args
     func_call_jaxpr = get_call_jaxpr(jaxpr)

frontend/test/pytest/test_debug.py

@@ -558,6 +558,12 @@ def test_no_options_to_mlir_opt(self):
         ).strip()
         assert expected in observed
 
+    def test_catalyst_error(self):
+        mlir = """This is invalid MLIR"""
+        msg = "custom op 'This'"
+        with pytest.raises(CompileError, match=msg):
+            to_mlir_opt(stdin=mlir)
+
 
 if __name__ == "__main__":
     pytest.main(["-x", __file__])

frontend/test/pytest/test_gradient.py

@@ -2201,5 +2201,78 @@ def circuit(x: float):
                 return qml.expval(qml.PauliZ(wires=0))
 
 
+def test_closure_variable_grad():
+    """Test that grad can take closure variables"""
+
+    @qml.qjit
+    def workflow_closure(x, y):
+
+        dev = qml.device("lightning.qubit", wires=1)
+
+        @qml.qnode(dev)
+        def circuit(x):
+            qml.RX(jnp.pi * x, wires=0)
+            qml.RX(jnp.pi * y, wires=0)
+            return qml.expval(qml.PauliY(0))
+
+        g = grad(circuit)
+        return g(x)
+
+    @qml.qjit
+    def workflow_no_closure(x, y):
+
+        dev = qml.device("lightning.qubit", wires=1)
+
+        @qml.qnode(dev)
+        def circuit(x, y):
+            qml.RX(jnp.pi * x, wires=0)
+            qml.RX(jnp.pi * y, wires=0)
+            return qml.expval(qml.PauliY(0))
+
+        g = grad(circuit)
+        return g(x, y)
+
+    expected = workflow_no_closure(1.0, 0.25)
+    observed = workflow_closure(1.0, 0.25)
+    assert np.allclose(expected, observed)
+
+
+def test_closure_variable_value_and_grad():
+    """Test that value and grad can take closure variables"""
+
+    @qml.qjit
+    def workflow_closure(x, y):
+
+        dev = qml.device("lightning.qubit", wires=1)
+
+        @qml.qnode(dev)
+        def circuit(x):
+            qml.RX(jnp.pi * x, wires=0)
+            qml.RX(jnp.pi * y, wires=0)
+            return qml.expval(qml.PauliY(0))
+
+        g = value_and_grad(circuit)
+        return g(x)
+
+    @qml.qjit
+    def workflow_no_closure(x, y):
+
+        dev = qml.device("lightning.qubit", wires=1)
+
+        @qml.qnode(dev)
+        def circuit(x, y):
+            qml.RX(jnp.pi * x, wires=0)
+            qml.RX(jnp.pi * y, wires=0)
+            return qml.expval(qml.PauliY(0))
+
+        g = value_and_grad(circuit)
+        return g(x, y)
+
+    x, y = 1.0, 0.25
+    expected = workflow_no_closure(x, y)
+    observed = workflow_closure(x, y)
+    assert np.allclose(expected, observed)
+
+
 if __name__ == "__main__":
     pytest.main(["-x", __file__])