Merge branch 'main' into bug-cond-args

doc/releases/changelog-dev.md

@@ -53,7 +53,8 @@
   Array([[1], [0], [1], [1], [0], [1],[0]], dtype=int64))
   ```
 
-* A new function `catalyst.passes.pipeline` allows the quantum circuit transformation pass pipeline for QNodes within a qjit-compiled workflow to be configured.
+* A new function `catalyst.passes.pipeline` allows the quantum circuit transformation pass pipeline
+  for QNodes within a qjit-compiled workflow to be configured.
   [(#1131)](https://github.com/PennyLaneAI/catalyst/pull/1131)
 
   ```python
@@ -92,10 +93,12 @@
         return jnp.abs(circuit_pipeline(x) - circuit_other(x))
   ```
 
-  For a list of available passes, please see the [catalyst.passes module documentation](https://docs.pennylane.ai/projects/catalyst/en/stable/code/__init__.html#module-catalyst.passes).
+  For a list of available passes, please see the
+  [catalyst.passes module documentation](https://docs.pennylane.ai/projects/catalyst/en/stable/code/__init__.html#module-catalyst.passes).
 
-  The pass pipeline order and options can be configured *globally* for a
-  qjit-compiled function, by using the `circuit_transform_pipeline` argument of the :func:`~.qjit` decorator.
+  The pass pipeline order and options can be configured *globally* for a qjit-compiled
+  function, by using the `circuit_transform_pipeline` argument of the :func:`~.qjit`
+  decorator.
 
   ```python
     my_passes = {
@@ -114,11 +117,13 @@
   Available MLIR passes are now documented and available within the
   [catalyst.passes module documentation](https://docs.pennylane.ai/projects/catalyst/en/stable/code/__init__.html#module-catalyst.passes).
 
-* Catalyst Autograph now supports updating a single index or a slice of JAX arrays using Python's array assignment operator syntax.
+* Catalyst Autograph now supports updating a single index or a slice of JAX arrays using Python's
+  array assignment operator syntax.
   [(#769)](https://github.com/PennyLaneAI/catalyst/pull/769)
   [(#1143)](https://github.com/PennyLaneAI/catalyst/pull/1143)
 
-  Using operator assignment syntax in favor of `at...op` expressions is now possible for the following operations:
+  Using operator assignment syntax in favor of `at...op` expressions is now possible for the
+  following operations:
   * `x[i] += y` in favor of `x.at[i].add(y)`
   * `x[i] -= y` in favor of `x.at[i].add(-y)`
   * `x[i] *= y` in favor of `x.at[i].multiply(y)`
@@ -144,8 +149,24 @@
 
 <h3>Improvements</h3>
 
-* Bufferization of `gradient.ForwardOp` and `gradient.ReverseOp` now requires 3 steps: `gradient-preprocessing`,
-  `gradient-bufferize`, and `gradient-postprocessing`. `gradient-bufferize` has a new rewrite for `gradient.ReturnOp`.
+* Bufferization of `gradient.ForwardOp` and `gradient.ReverseOp` now requires 3 steps:
+  `gradient-preprocessing`, `gradient-bufferize`, and `gradient-postprocessing`.
+  `gradient-bufferize` has a new rewrite for `gradient.ReturnOp`.
+  [(#1139)](https://github.com/PennyLaneAI/catalyst/pull/1139)
+
+* Scalar tensors are eliminated from control flow operations in the program, and are replaced with
+  bare scalars instead. This improves compilation time and memory usage at runtime by avoiding heap
+  allocations and reducing the amount of instructions.
+  [(#1075)](https://github.com/PennyLaneAI/catalyst/pull/1075)
+
+  A new MLIR pass `detensorize-scf` is added that works in conjuction with the existing
+  `linalg-detensorize` pass to detensorize input programs. The IR generated by JAX wraps all values
+  in the program in tensors, including scalars, leading to unnecessary memory allocations for
+  programs compiled to CPU via MLIR to LLVM pipeline.
+
+* Bufferization of `gradient.ForwardOp` and `gradient.ReverseOp` now requires 3 steps:
+  `gradient-preprocessing`, `gradient-bufferize`, and `gradient-postprocessing`.
+  `gradient-bufferize` has a new rewrite for `gradient.ReturnOp`.
   [(#1139)](https://github.com/PennyLaneAI/catalyst/pull/1139)
 
 * The decorator `self_inverses` now supports all Hermitian Gates.
@@ -168,11 +189,10 @@
   Three-bit Gates: Toffoli
   - [`qml.Toffoli`](https://docs.pennylane.ai/en/stable/code/api/pennylane.Toffoli.html)
 
-
-
 * Support is expanded for backend devices that exculsively return samples in the measurement
   basis. Pre- and post-processing now allows `qjit` to be used on these devices with `qml.expval`,
-  `qml.var` and `qml.probs` measurements in addiiton to `qml.sample`, using the `measurements_from_samples` transform.
+  `qml.var` and `qml.probs` measurements in addiiton to `qml.sample`, using the
+  `measurements_from_samples` transform.
   [(#1106)](https://github.com/PennyLaneAI/catalyst/pull/1106)
 
 * Catalyst now supports numpy 2.0
@@ -187,6 +207,9 @@
 
 <h3>Bug fixes</h3>
 
+* Resolve a bug in the `vmap` function when passing shapeless values to the target.
+  [(#1150)](https://github.com/PennyLaneAI/catalyst/pull/1150)
+
 * Fix error message displayed when using `qml.cond` on callables with arguments.
   [(#1151)](https://github.com/PennyLaneAI/catalyst/pull/1151)
 
@@ -199,8 +222,9 @@
   running gradient like functions.
   [(#1063)](https://github.com/PennyLaneAI/catalyst/pull/1063)
 
-* Functions with multiple tapes are now split with a new mlir pass `--split-multiple-tapes`, with one tape per function.
-  The reset routine that makes a maeasurement between tapes and inserts a X gate if measured one is no longer used.
+* Functions with multiple tapes are now split with a new mlir pass `--split-multiple-tapes`, with
+  one tape per function. The reset routine that makes a maeasurement between tapes and inserts a
+  X gate if measured one is no longer used.
   [(#1017)](https://github.com/PennyLaneAI/catalyst/pull/1017)
   [(#1130)](https://github.com/PennyLaneAI/catalyst/pull/1130)
 
@@ -215,17 +239,26 @@
 * The device capability loading mechanism has been moved into the `QJITDevice` constructor.
   [(#1141)](https://github.com/PennyLaneAI/catalyst/pull/1141)
 
+* Several functions related to device capabilities have been refactored.
+  [(#1149)](https://github.com/PennyLaneAI/catalyst/pull/1149)
+
+  In particular, the signatures of `get_device_capability`, `catalyst_decompose`,
+ `catalyst_acceptance`, and `QJITDevice.__init__` have changed, and the `pennylane_operation_set`
+  function has been removed entirely.
+
 <h3>Contributors</h3>
 
 This release contains contributions from (in alphabetical order):
 
 Joey Carter,
+Spencer Comin,
 Lillian M.A. Frederiksen,
+Sengthai Heng,
+David Ittah,
+Mehrdad Malekmohammadi,
+Vincent Michaud-Rioux,
 Romain Moyard,
 Erick Ochoa Lopez,
-Mehrdad Malekmohammadi,
-Paul Haochen Wang,
-Sengthai Heng,
-Spencer Comin,
 Daniel Strano,
-Raul Torres.
+Raul Torres,
+Paul Haochen Wang.

frontend/catalyst/_version.py

@@ -16,4 +16,4 @@
    Version number (major.minor.patch[-label])
 """
 
-__version__ = "0.9.0-dev20"
+__version__ = "0.9.0-dev21"

frontend/catalyst/api_extensions/function_maps.py

@@ -319,10 +319,18 @@ def _get_batch_size(self, args_flat, axes_flat, axis_size):
         """
 
         batch_sizes = []
-        for arg, d in zip(args_flat, axes_flat):
-            shape = np.shape(arg) if arg.shape else (1,)
-            if d is not None and len(shape) > d:
+        for i, (arg, d) in enumerate(zip(args_flat, axes_flat)):
+            if d is None:
+                continue
+
+            shape = np.shape(arg)
+            if len(shape) > d:
                 batch_sizes.append(shape[d])
+            else:
+                raise ValueError(
+                    f"Axis specifier {d} is out of bounds for argument {i}. "
+                    f"The argument only has {len(shape)} dimensions."
+                )
 
         if any(size != batch_sizes[0] for size in batch_sizes[1:]):
             raise ValueError(

frontend/catalyst/compiler.py

@@ -197,6 +197,8 @@ def run_writing_command(command: List[str], compile_options: Optional[CompileOpt
         "hlo-custom-call-lowering",
         "cse",
         "func.func(linalg-detensorize{aggressive-mode})",
+        "detensorize-scf",
+        "canonicalize",
     ],
 )
 

frontend/catalyst/device/decomposition.py

@@ -49,10 +49,8 @@ def check_alternative_support(op, capabilities):
 
     if isinstance(op, qml.ops.Controlled):
         # "Cast" away the specialized class for gates like Toffoli, ControlledQubitUnitary, etc.
-        if (
-            capabilities.native_ops.get(op.base.name)
-            and capabilities.native_ops.get(op.base.name).controllable
-        ):
+        supported = capabilities.native_ops.get(op.base.name)
+        if supported and supported.controllable:
             return [qml.ops.Controlled(op.base, op.control_wires, op.control_values, op.work_wires)]
 
     return None
@@ -71,7 +69,7 @@ def catalyst_decomposer(op, capabilities: DeviceCapabilities):
 
     if capabilities.native_ops.get("QubitUnitary"):
         # If the device supports unitary matrices, apply the relevant conversions and fallbacks.
-        if capabilities.to_matrix_ops.get(op.name) or (
+        if op.name in capabilities.to_matrix_ops or (
             op.has_matrix and isinstance(op, qml.ops.Controlled)
         ):
             return _decompose_to_matrix(op)
@@ -81,13 +79,7 @@ def catalyst_decomposer(op, capabilities: DeviceCapabilities):
 
 @transform
 @debug_logger
-def catalyst_decompose(
-    tape: qml.tape.QuantumTape,
-    ctx,
-    stopping_condition,
-    capabilities,
-    max_expansion=None,
-):
+def catalyst_decompose(tape: qml.tape.QuantumTape, ctx, capabilities):
     """Decompose operations until the stopping condition is met.
 
     In a single call of the catalyst_decompose function, the PennyLane operations are decomposed
@@ -104,18 +96,17 @@ def catalyst_decompose(
 
     (toplevel_tape,), _ = decompose(
         tape,
-        stopping_condition,
+        stopping_condition=lambda op: bool(catalyst_acceptance(op, capabilities)),
         skip_initial_state_prep=capabilities.initial_state_prep_flag,
         decomposer=partial(catalyst_decomposer, capabilities=capabilities),
-        max_expansion=max_expansion,
         name="catalyst on this device",
         error=CompileError,
     )
 
     new_ops = []
     for op in toplevel_tape.operations:
         if has_nested_tapes(op):
-            op = _decompose_nested_tapes(op, ctx, stopping_condition, capabilities, max_expansion)
+            op = _decompose_nested_tapes(op, ctx, capabilities)
         new_ops.append(op)
     tape = qml.tape.QuantumScript(new_ops, tape.measurements, shots=tape.shots)
 
@@ -133,19 +124,15 @@ def _decompose_to_matrix(op):
     return [op]
 
 
-def _decompose_nested_tapes(op, ctx, stopping_condition, capabilities, max_expansion):
+def _decompose_nested_tapes(op, ctx, capabilities):
     new_regions = []
     for region in op.regions:
         if region.quantum_tape is None:
             new_tape = None
         else:
             with EvaluationContext.frame_tracing_context(ctx, region.trace):
                 tapes, _ = catalyst_decompose(
-                    region.quantum_tape,
-                    ctx=ctx,
-                    stopping_condition=stopping_condition,
-                    capabilities=capabilities,
-                    max_expansion=max_expansion,
+                    region.quantum_tape, ctx=ctx, capabilities=capabilities
                 )
                 new_tape = tapes[0]
         new_regions.append(
@@ -198,9 +185,22 @@ def null_postprocessing(results):
     return [new_tape], null_postprocessing
 
 
-def catalyst_acceptance(op: qml.operation.Operator, operations) -> bool:
-    """Specify whether or not an Operator is supported."""
-    return op.name in operations
+def catalyst_acceptance(op: qml.operation.Operation, capabilities: DeviceCapabilities) -> str:
+    """Check whether or not an Operator is supported."""
+    op_support = capabilities.native_ops
+
+    if isinstance(op, qml.ops.Adjoint):
+        match = catalyst_acceptance(op.base, capabilities)
+        if not match or not op_support[match].invertible:
+            return None
+    elif type(op) is qml.ops.ControlledOp:
+        match = catalyst_acceptance(op.base, capabilities)
+        if not match or not op_support[match].controllable:
+            return None
+    else:
+        match = op.name if op.name in op_support else None
+
+    return match
 
 
 @transform