Implement automatic qubit management

doc/releases/changelog-dev.md

@@ -74,6 +74,38 @@
   Array([0.25, 0.25, 0.25, 0.25], dtype=float64))
   ```
 
+* Catalyst now supports automatic qubit management.
+  [(#1788)](https://github.com/PennyLaneAI/catalyst/pull/1788)
+
+  The number of wires does not need to be speficied during device initialization,
+  and instead will be automatically managed by the Catalyst Runtime.
+
+  ```python
+  @qjit
+  def workflow():
+      dev = qml.device("lightning.qubit") # no wires here!
+      @qml.qnode(dev)
+      def circuit():
+          qml.PauliX(wires=2)
+          return qml.probs()
+      return circuit()
+
+  print(workflow())
+  ```
+
+  ```pycon
+  [0. 1. 0. 0. 0. 0. 0. 0.]
+  ```
+
+  In this example, the number of wires is not specified at device initialization.
+  When we encounter an X gate on `wires=2`, catalyst automatically expands the size
+  of the qubit register to include the requested wire index.
+  Here, the register will contain (at least) 3 qubits after the X operation.
+  As a result, we can see the QNode returning the probabilities for the state |001>,
+  meaning 3 wires were allocated in total.
+
+  This feature can be turned on by omitting the `wires` argument to the device.
+
 <h3>Improvements 🛠</h3>
 
 * The behaviour of measurement processes executed on `null.qubit` with QJIT is now more in line with

frontend/catalyst/device/qjit_device.py

@@ -548,14 +548,17 @@ def is_dynamic_wires(wires: qml.wires.Wires):
     If the number of wires is dynamic, the Wires object contains a single tracer that
     represents the number of wires.
     """
+    # Automatic qubit management mode should not encounter this query
+    assert wires is not None
     return (len(wires) == 1) and (isinstance(wires[0], DynamicJaxprTracer))
 
 
 def check_device_wires(wires):
     """Validate requirements Catalyst imposes on device wires."""
 
     if wires is None:
-        raise AttributeError("Catalyst does not support device instances without set wires.")
+        # Automatic qubit management mode, nothing to check
+        return
 
     if len(wires) >= 2 or (not is_dynamic_wires(wires)):
         # A dynamic number of wires correspond to a single tracer for the number

frontend/catalyst/from_plxpr.py

@@ -301,7 +301,11 @@ def __setattr__(self, __name: str, __value) -> None:
     def setup(self):
         """Initialize the stateref and bind the device."""
         if self.stateref is None:
-            device_init_p.bind(self._shots, **_get_device_kwargs(self._device))
+            device_init_p.bind(
+                self._shots,
+                auto_qubit_management=(self._device.wires is None),
+                **_get_device_kwargs(self._device),
+            )
             self.stateref = {"qreg": qalloc_p.bind(len(self._device.wires)), "wire_map": {}}
 
     # pylint: disable=attribute-defined-outside-init

frontend/catalyst/jax_primitives.py

@@ -813,12 +813,18 @@ def _zne_lowering(ctx, *args, folding, jaxpr, fn):
 # device_init
 #
 @device_init_p.def_abstract_eval
-def _device_init_abstract_eval(shots, rtd_lib, rtd_name, rtd_kwargs):
+def _device_init_abstract_eval(shots, auto_qubit_management, rtd_lib, rtd_name, rtd_kwargs):
     return ()
 
 
+# pylint: disable=too-many-arguments, too-many-positional-arguments
 def _device_init_lowering(
-    jax_ctx: mlir.LoweringRuleContext, shots: ir.Value, rtd_lib, rtd_name, rtd_kwargs
+    jax_ctx: mlir.LoweringRuleContext,
+    shots: ir.Value,
+    auto_qubit_management,
+    rtd_lib,
+    rtd_name,
+    rtd_kwargs,
 ):
     ctx = jax_ctx.module_context.context
     ctx.allow_unregistered_dialects = True
@@ -829,6 +835,7 @@ def _device_init_lowering(
         ir.StringAttr.get(rtd_name),
         ir.StringAttr.get(rtd_kwargs),
         shots=shots_value,
+        auto_qubit_management=auto_qubit_management,
     )
 
     return ()

frontend/catalyst/jax_tracer.py

@@ -985,11 +985,13 @@ def trace_quantum_measurements(
                     "Use qml.sample() instead."
                 )
 
-            d_wires = (
-                num_qubits_p.bind()
-                if catalyst.device.qjit_device.is_dynamic_wires(device.wires)
-                else len(device.wires)
-            )
+            if device.wires is None:
+                d_wires = num_qubits_p.bind()
+            elif catalyst.device.qjit_device.is_dynamic_wires(device.wires):
+                d_wires = num_qubits_p.bind()
+            else:
+                d_wires = len(device.wires)
+
             m_wires = output.wires if output.wires else None
             obs_tracers, nqubits = trace_observables(output.obs, qrp, m_wires)
             nqubits = d_wires if nqubits is None else nqubits
@@ -1386,11 +1388,16 @@ def is_leaf(obj):
                 device_shots = get_device_shots(device) or 0
                 device_init_p.bind(
                     device_shots,
+                    auto_qubit_management=(device.wires is None),
                     rtd_lib=device.backend_lib,
                     rtd_name=device.backend_name,
                     rtd_kwargs=str(device.backend_kwargs),
                 )
-                if catalyst.device.qjit_device.is_dynamic_wires(device.wires):
+                if device.wires is None:
+                    # Automatic qubit management mode
+                    # We start with 0 wires and allocate new wires in runtime as we encounter them.
+                    qreg_in = qalloc_p.bind(0)
+                elif catalyst.device.qjit_device.is_dynamic_wires(device.wires):
                     # When device has dynamic wires, the device.wires iterable object
                     # has a single value, which is the tracer for the number of wires
                     qreg_in = qalloc_p.bind(device.wires[0])

frontend/test/lit/test_measurements.py

@@ -678,3 +678,26 @@ def circ():
 
 state_dynamic(10)
 print(state_dynamic.mlir)
+
+
+# CHECK-LABEL: @automatic_qubit_management
+@qjit(target="mlir")
+def automatic_qubit_management():
+    @qml.qnode(qml.device("lightning.qubit"))
+    def circ():
+        # CHECK: [[qreg:%.+]] = quantum.alloc( 0) : !quantum.reg
+        # CHECK: [[in_qubit:%.+]] = quantum.extract [[qreg]][ 2] : !quantum.reg -> !quantum.bit
+        # CHECK: [[out_qubit:%.+]] = quantum.custom "Hadamard"() [[in_qubit]] : !quantum.bit
+        qml.Hadamard(wires=2)
+
+        # CHECK: [[nqubits:%.+]] = quantum.num_qubits : i64
+        # CHECK: [[toTensor:%.+]] = tensor.from_elements [[nqubits]] : tensor<i64>
+        # CHECK: [[probs_shape:%.+]] = stablehlo.shift_left {{%.+}}, [[toTensor]] : tensor<i64>
+        # CHECK: [[deTensor:%.+]] = tensor.extract [[probs_shape]][] : tensor<i64>
+        # CHECK: {{%.+}} = quantum.probs {{%.+}} shape [[deTensor]] : tensor<?xf64>
+        return qml.probs()
+
+    return circ()
+
+
+print(automatic_qubit_management.mlir)

frontend/test/pytest/test_automatic_qubit_management.py

@@ -0,0 +1,153 @@
+# Copyright 2025 Xanadu Quantum Technologies Inc.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+This file contains tests for automatic qubit management.
+Automatic qubit management refers to when the user does not specify the total number of wires
+during device initialization.
+
+Note that, catalyst and pennylane handles device labels differently. For example, when a new label
+    qml.gate_or_measurement(wires=1000)
+is encountered, core pennylane considers "1000" as a pure wire label, and interprets that as
+*one* new wire, with the label "1000". However in catalyst we do not associate wires with
+arbitrary labels and require wires to be continuous integers from zero, and we would interpret this
+as "allocate new wires until we have 1001 wires, and act on wire[1000]".
+
+In other words, the reference runs of automatic qubit management should be qjit runs with wires
+specified during device initialization, instead of non qjit runs.
+"""
+
+import numpy as np
+import pennylane as qml
+
+from catalyst import qjit
+
+
+def test_partial_sample(backend):
+    """
+    Test that a `sample` terminal measurement with wires specified can be executed
+    correctly with automatic qubit management.
+    """
+
+    def circuit():
+        qml.RX(0.0, wires=0)
+        return qml.sample(wires=[0, 2])
+
+    wires = [4, None]
+    devices = [qml.device(backend, shots=10, wires=wire) for wire in wires]
+    ref, observed = (qjit(qml.qnode(dev)(circuit))() for dev in devices)
+    assert ref.shape == observed.shape
+    assert np.allclose(ref, observed)
+
+
+def test_partial_counts(backend):
+    """
+    Test that a `counts` terminal measurement with wires specified can be executed
+    correctly with automatic qubit management.
+    """
+
+    def circuit():
+        qml.RX(0.0, wires=0)
+        return qml.counts(wires=[0, 2])
+
+    wires = [4, None]
+    devices = [qml.device(backend, shots=10, wires=wire) for wire in wires]
+    ref, observed = (qjit(qml.qnode(dev)(circuit))() for dev in devices)
+    assert (ref[i].shape == observed[i].shape for i in (0, 1))
+    assert np.allclose(ref, observed)
+
+
+def test_partial_probs(backend):
+    """
+    Test that a `probs` terminal measurement with wires specified can be executed
+    correctly with automatic qubit management.
+    """
+
+    def circuit():
+        qml.PauliX(wires=0)
+        return qml.probs(wires=[0, 2])
+
+    wires = [4, None]
+    devices = [qml.device(backend, wires=wire) for wire in wires]
+    ref, observed = (qjit(qml.qnode(dev)(circuit))() for dev in devices)
+    assert ref.shape == observed.shape
+    assert np.allclose(ref, observed)
+
+
+def test_sample(backend):
+    """
+    Test that a `sample` terminal measurement without wires specified can be executed
+    correctly with automatic qubit management.
+    """
+
+    def circuit():
+        qml.RX(0.0, wires=3)
+        return qml.sample()
+
+    wires = [4, None]
+    devices = [qml.device(backend, shots=10, wires=wire) for wire in wires]
+    ref, observed = (qjit(qml.qnode(dev)(circuit))() for dev in devices)
+    assert ref.shape == observed.shape
+    assert np.allclose(ref, observed)
+
+
+def test_counts(backend):
+    """
+    Test that a `counts` terminal measurement without wires specified can be executed
+    correctly with automatic qubit management.
+    """
+
+    def circuit():
+        qml.RX(0.0, wires=3)
+        return qml.counts()
+
+    wires = [4, None]
+    devices = [qml.device(backend, shots=10, wires=wire) for wire in wires]
+    ref, observed = (qjit(qml.qnode(dev)(circuit))() for dev in devices)
+    assert (ref[i].shape == observed[i].shape for i in (0, 1))
+    assert np.allclose(ref, observed)
+
+
+def test_probs(backend):
+    """
+    Test that a `probs` terminal measurement without wires specified can be executed
+    correctly with automatic qubit management.
+    """
+
+    def circuit():
+        qml.PauliX(wires=3)
+        return qml.probs()
+
+    wires = [4, None]
+    devices = [qml.device(backend, wires=wire) for wire in wires]
+    ref, observed = (qjit(qml.qnode(dev)(circuit))() for dev in devices)
+    assert ref.shape == observed.shape
+    assert np.allclose(ref, observed)
+
+
+def test_state(backend):
+    """
+    Test that a `state` terminal measurement without wires specified can be executed
+    correctly with automatic qubit management.
+    """
+
+    def circuit():
+        qml.PauliX(wires=3)
+        return qml.state()
+
+    wires = [4, None]
+    devices = [qml.device(backend, wires=wire) for wire in wires]
+    ref, observed = (qjit(qml.qnode(dev)(circuit))() for dev in devices)
+    assert ref.shape == observed.shape
+    assert np.allclose(ref, observed)

frontend/test/pytest/test_device_api.py

@@ -56,17 +56,6 @@ def test_qjit_device():
         device_qjit.execute(10, 2)
 
 
-def test_qjit_device_no_wires():
-    """Test the qjit device from a device using the new api without wires set."""
-    device = NullQubit(shots=2032)
-
-    with pytest.raises(
-        AttributeError, match="Catalyst does not support device instances without set wires."
-    ):
-        # Create qjit device
-        QJITDevice(device)
-
-
 @pytest.mark.parametrize(
     "wires",
     (

mlir/include/Quantum/IR/QuantumOps.td

@@ -81,6 +81,7 @@ def DeviceInitOp : Quantum_Op<"device"> {
 
     let arguments = (ins
         Optional<I64>:$shots,
+        UnitAttr:$auto_qubit_management,
         StrAttr:$lib,
         StrAttr:$device_name,
         StrAttr:$kwargs
@@ -90,6 +91,18 @@ def DeviceInitOp : Quantum_Op<"device"> {
       (`shots` `(` $shots^ `)`)? `[` $lib `,` $device_name `,` $kwargs `]` attr-dict
     }];
 
+    let builders = [
+        OpBuilder<
+        // Convenience builder for a device not in automatic qubit management mode
+        (ins
+             "mlir::Value":$shots,
+             "mlir::StringAttr":$lib,
+             "mlir::StringAttr":$device_name,
+             "mlir::StringAttr":$kwargs
+        ),[{
+            DeviceInitOp::build($_builder, $_state, shots, false, lib, device_name, kwargs);
+        }]>,
+    ];
 }
 
 def DeviceReleaseOp : Quantum_Op<"device_release"> {

mlir/lib/Quantum/Transforms/ConversionPatterns.cpp

@@ -187,11 +187,13 @@ struct DeviceInitOpPattern : public OpConversionPattern<DeviceInitOp> {
 
         Type charPtrType = LLVM::LLVMPointerType::get(rewriter.getContext());
         Type int64Type = IntegerType::get(rewriter.getContext(), 64);
+        Type int1Type = IntegerType::get(rewriter.getContext(), 1);
         Type qirSignature = LLVM::LLVMFunctionType::get(LLVM::LLVMVoidType::get(ctx),
                                                         {/* rtd_lib = */ charPtrType,
                                                          /* rtd_name = */ charPtrType,
                                                          /* rtd_kwargs = */ charPtrType,
-                                                         /* shots = */ int64Type});
+                                                         /* shots = */ int64Type,
+                                                         /*auto_qubit_management = */ int1Type});
         LLVM::LLVMFuncOp fnDecl =
             catalyst::ensureFunctionDeclaration(rewriter, op, qirName, qirSignature);
 
@@ -217,6 +219,10 @@ struct DeviceInitOpPattern : public OpConversionPattern<DeviceInitOp> {
             operands.push_back(shots);
         }
 
+        Value autoQubitManagement = rewriter.create<LLVM::ConstantOp>(loc, rewriter.getI1Type(),
+                                                                      op.getAutoQubitManagement());
+        operands.push_back(autoQubitManagement);
+
         rewriter.create<LLVM::CallOp>(loc, fnDecl, operands);
 
         rewriter.eraseOp(op);