Implement automatic qubit management (#1788)

**Context:**
We wish to support automatic qubit management, i.e. when the user does
not specify the number of wires in device initialization
```python
@qjit
def workflow():
    dev = qml.device("lightning.qubit") # no wires here!
    @qml.qnode(dev)
    def circuit():
        qml.PauliX(wires=0)
        return qml.probs(wires=[2,3])  # non qjit is [1,0,0,0]
    return circuit()
```

**Description of the Change:**
A new unit attribute, `auto_qubit_management`, is added to the device
init op. If no wires are specified during device init in user python, we
trace to device init op with this attribute turned on.

In runtime, when a new user wire index is encountered in
`__catalyst__rt__array_get_element_ptr_1d`, if the device is in
automatic management mode (as indicated by the `RTDevice` constructor,
set from the new attribute we added), we allocate new device qubits so
that this new user wire index is an available qubit.

**Benefits:**
The automatic qubit management feature is now supported.

**Possible Drawbacks:**
1. Catalyst and Pennylane interprets wire labels differently. See #1800 
2. Potential out-of-memory error might be encountered in runtime.

[sc-86500]

Author: paul0403

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);