Handle missing cases in ObserveAnsatz

lib/Optimizer/Transforms/ObserveAnsatz.cpp

@@ -77,6 +77,9 @@ struct AnsatzMetadata {
 
   /// Map qubit indices to their mlir Value
   DenseMap<std::size_t, Value> qubitValues;
+
+  /// Map qubit indices to their mlir veq Value and index
+  DenseMap<std::size_t, std::pair<Value, std::size_t>> qubitValuesIndexed;
 };
 
 /// Define a map type for Quake functions to their associated metadata
@@ -140,17 +143,20 @@ struct AnsatzFunctionAnalysis {
 
     // walk and find all quantum allocations
     auto walkResult = funcOp->walk([&](quake::AllocaOp op) {
-      if (auto veq = dyn_cast<quake::VeqType>(op.getResult().getType())) {
-        // Only update data.nQubits here. data.qubitValues will be updated for
-        // the corresponding ExtractRefOP's in the `walk` below.
-        if (veq.hasSpecifiedSize())
-          data.nQubits += veq.getSize();
-        else
+      Value result = op.getResult();
+      if (auto veq = dyn_cast<quake::VeqType>(result.getType())) {
+        // Update data.nQubits here and store qubit info as indexes into
+        // the veq, in case we don't encounter any ExtractRefOps for
+        // them later.
+        if (veq.hasSpecifiedSize()) {
+          for (std::size_t i = 0; i < veq.getSize(); i++)
+            data.qubitValuesIndexed.insert({data.nQubits++, {result, i}});
+        } else
           return WalkResult::interrupt(); // this is an error condition
       } else {
         // single alloc is for a single qubit. Update data.qubitValues here
         // because ExtractRefOp `walk` won't find any ExtractRefOp for this.
-        data.qubitValues.insert({data.nQubits++, op.getResult()});
+        data.qubitValues.insert({data.nQubits++, result});
       }
       return WalkResult::advance();
     });
@@ -259,7 +265,8 @@ struct AppendMeasurements : public OpRewritePattern<func::FuncOp> {
     auto loc = funcOp.getBody().back().getTerminator()->getLoc();
     Operation *last = &funcOp.getBody().back().front();
     funcOp.walk([&](Operation *op) {
-      if (dyn_cast<quake::OperatorInterface>(op))
+      if (dyn_cast<quake::OperatorInterface>(op) ||
+          dyn_cast<quake::AllocaOp>(op))
         last = op;
     });
     builder.setInsertionPointAfter(last);
@@ -279,10 +286,20 @@ struct AppendMeasurements : public OpRewritePattern<func::FuncOp> {
       // do nothing for z or identities
 
       // get the qubit value
+      Value qubitVal;
       auto seek = iter->second.qubitValues.find(i);
-      if (seek == iter->second.qubitValues.end())
-        continue;
-      auto qubitVal = seek->second;
+      if (seek == iter->second.qubitValues.end()) {
+        auto seekIndexed = iter->second.qubitValuesIndexed.find(i);
+        if (seekIndexed == iter->second.qubitValuesIndexed.end())
+          continue;
+        auto veqOp = seekIndexed->second.first;
+        auto index = seekIndexed->second.second;
+        auto extractRef =
+            builder.create<quake::ExtractRefOp>(loc, veqOp, index);
+        qubitVal = extractRef.getResult();
+      } else {
+        qubitVal = seek->second;
+      }
 
       // append the measurement basis change ops
       // Note: when using value semantics, qubitVal will be updated to the new

python/tests/backends/test_Ionq_LocalEmulation_kernel.py

@@ -8,11 +8,16 @@
 
 import cudaq
 import cudaq.kernels
+from cudaq import spin
 import pytest
 import os
 from typing import List
 
 
+def assert_close(want, got, tolerance=1.0e-1) -> bool:
+    return abs(want - got) < tolerance
+
+
 @pytest.fixture(scope="function", autouse=True)
 def configureTarget():
     # Set the targeted QPU
@@ -23,6 +28,19 @@ def configureTarget():
     cudaq.reset_target()
 
 
+def test_Ionq_observe():
+    cudaq.set_random_seed(13)
+
+    @cudaq.kernel
+    def ansatz():
+        q = cudaq.qvector(1)
+
+    molecule = 5.0 - 1.0 * spin.x(0)
+    res = cudaq.observe(ansatz, molecule, shots_count=10000)
+    print(res.expectation())
+    assert assert_close(5.0, res.expectation())
+
+
 def test_Ionq_cudaq_uccsd():
 
     num_electrons = 2

python/tests/backends/test_Quantinuum_LocalEmulation_kernel.py

@@ -14,8 +14,8 @@
 from typing import List
 
 
-def assert_close(got) -> bool:
-    return got < -1.5 and got > -1.9
+def assert_close(want, got, tolerance=1.0e-1) -> bool:
+    return abs(want - got) < tolerance
 
 
 @pytest.fixture(scope="function", autouse=True)
@@ -85,13 +85,26 @@ def ansatz(theta: float):
 
     # Run the observe task on quantinuum synchronously
     res = cudaq.observe(ansatz, hamiltonian, .59, shots_count=100000)
-    assert assert_close(res.expectation())
+    assert assert_close(-1.7, res.expectation())
 
     # Launch it asynchronously, enters the job into the queue
     future = cudaq.observe_async(ansatz, hamiltonian, .59, shots_count=100000)
     # Retrieve the results (since we're emulating)
     res = future.get()
-    assert assert_close(res.expectation())
+    assert assert_close(-1.7, res.expectation())
+
+
+def test_observe():
+    cudaq.set_random_seed(13)
+
+    @cudaq.kernel
+    def ansatz():
+        q = cudaq.qvector(1)
+
+    molecule = 5.0 - 1.0 * spin.x(0)
+    res = cudaq.observe(ansatz, molecule, shots_count=10000)
+    print(res.expectation())
+    assert assert_close(5.0, res.expectation())
 
 
 def test_quantinuum_exp_pauli():
@@ -111,13 +124,13 @@ def ansatz(theta: float):
 
     # Run the observe task on quantinuum synchronously
     res = cudaq.observe(ansatz, hamiltonian, .59, shots_count=100000)
-    assert assert_close(res.expectation())
+    assert assert_close(-1.7, res.expectation())
 
     # Launch it asynchronously, enters the job into the queue
     future = cudaq.observe_async(ansatz, hamiltonian, .59, shots_count=100000)
     # Retrieve the results (since we're emulating)
     res = future.get()
-    assert assert_close(res.expectation())
+    assert assert_close(-1.7, res.expectation())
 
 
 def test_u3_emulatation():

python/tests/kernel/test_kernel_features.py

@@ -185,6 +185,17 @@ def grover(N: int, M: int, oracle: Callable[[cudaq.qview], None]):
     assert '011' in counts
 
 
+def test_observe():
+
+    @cudaq.kernel
+    def ansatz():
+        q = cudaq.qvector(1)
+
+    molecule = 5.0 - 1.0 * spin.x(0)
+    res = cudaq.observe(ansatz, molecule, shots_count=10000)
+    assert np.isclose(res.expectation(), 5.0, atol=1e-1)
+
+
 def test_pauli_word_input():
 
     h2_data = [

targettests/execution/cudaq_observe_qvector.cpp

@@ -0,0 +1,49 @@
+/*******************************************************************************
+ * Copyright (c) 2022 - 2025 NVIDIA Corporation & Affiliates.                  *
+ * All rights reserved.                                                        *
+ *                                                                             *
+ * This source code and the accompanying materials are made available under    *
+ * the terms of the Apache License 2.0 which accompanies this distribution.    *
+ ******************************************************************************/
+
+// REQUIRES: c++20
+// clang-format off
+// RUN: nvq++ %cpp_std --target infleqtion      --emulate %s -o %t && %t | FileCheck %s
+// RUN: nvq++ --target ionq                     --emulate %s -o %t && %t | FileCheck %s
+// 2 different IQM machines for 2 different topologies
+// RUN: nvq++ --target iqm --iqm-machine Adonis --emulate %s -o %t && %t | FileCheck %s
+// RUN: nvq++ --target iqm --iqm-machine Apollo --emulate %s -o %t && %t | FileCheck %s
+// RUN: nvq++ --target oqc                      --emulate %s -o %t && %t | FileCheck %s
+// RUN: nvq++ --target quantinuum               --emulate %s -o %t && %t | FileCheck %s
+// RUN: if %braket_avail; then nvq++ --target braket --emulate %s -o %t && %t | FileCheck %s; fi
+// clang-format on
+
+#include <cudaq.h>
+#include <cudaq/algorithm.h>
+
+// The example here shows a simple use case for the `cudaq::observe`
+// function in computing expected values of provided spin_ops.
+
+struct ansatz {
+  auto operator()() __qpu__ {
+    cudaq::qvector q(1);
+  }
+};
+
+int main() {
+
+  // Build up your spin op algebraically
+  using namespace cudaq::spin;
+  cudaq::spin_op h = 5.0 - 1.0 * x(0);
+
+  // Make repeatable for shots-based emulation
+  cudaq::set_random_seed(13);
+
+  // Observe takes the kernel, the spin_op, and the concrete
+  // parameters for the kernel
+  int energy = (int)(cudaq::observe(10000, ansatz{}, h) + 0.5);
+  printf("Energy is %d.\n", energy);
+  return 0;
+}
+
+// CHECK: Energy is 5.

test/Quake/observeAnsatz.qke

@@ -25,7 +25,58 @@
     return
   }
 
-// CHECK: quake.h %
-// CHECK: quake.h %
-// CHECK: %{{.*}} = quake.mz %{{.*}} : (!quake.ref) -> !quake.measure
-// CHECK: %{{.*}} = quake.mz %{{.*}} : (!quake.ref) -> !quake.measure
+// CHECK-LABEL:   func.func @__nvqpp__mlirgen__ansatz(%arg0: f64) {
+// CHECK:           quake.h %
+// CHECK:           quake.h %
+// CHECK:           %{{.*}} = quake.mz %{{.*}} : (!quake.ref) -> !quake.measure
+// CHECK:           %{{.*}} = quake.mz %{{.*}} : (!quake.ref) -> !quake.measure
+
+  func.func @test_veq_no_extract_refs() {
+    %0 = quake.alloca !quake.veq<2>
+    return
+  }
+
+// CHECK-LABEL:   func.func @test_veq_no_extract_refs() {
+// CHECK:           %[[VAL_0:.*]] = quake.alloca !quake.veq<2>
+// CHECK:           %[[VAL_1:.*]] = quake.extract_ref %[[VAL_0]][0] : (!quake.veq<2>) -> !quake.ref
+// CHECK:           quake.h %[[VAL_1]] : (!quake.ref) -> ()
+// CHECK:           %[[VAL_2:.*]] = quake.extract_ref %[[VAL_0]][1] : (!quake.veq<2>) -> !quake.ref
+// CHECK:           quake.h %[[VAL_2]] : (!quake.ref) -> ()
+// CHECK:           %[[VAL_3:.*]] = quake.mz %[[VAL_1]] name "r00000" : (!quake.ref) -> !quake.measure
+// CHECK:           %[[VAL_4:.*]] = quake.mz %[[VAL_2]] name "r00001" : (!quake.ref) -> !quake.measure
+// CHECK:           return
+// CHECK:         }
+
+  func.func @test_veq_no_extract_refs2() {
+    %0 = quake.alloca !quake.ref
+    %1 = quake.alloca !quake.ref
+    return
+  }
+
+// CHECK-LABEL:   func.func @test_veq_no_extract_refs2() {
+// CHECK:           %[[VAL_0:.*]] = quake.alloca !quake.ref
+// CHECK:           %[[VAL_1:.*]] = quake.alloca !quake.ref
+// CHECK:           quake.h %[[VAL_0]] : (!quake.ref) -> ()
+// CHECK:           quake.h %[[VAL_1]] : (!quake.ref) -> ()
+// CHECK:           %[[VAL_3:.*]] = quake.mz %[[VAL_0]] name "r00000" : (!quake.ref) -> !quake.measure
+// CHECK:           %[[VAL_4:.*]] = quake.mz %[[VAL_1]] name "r00001" : (!quake.ref) -> !quake.measure
+// CHECK:           return
+// CHECK:         }
+
+  func.func @test_veq_no_extract_refs3() {
+    %0 = quake.alloca !quake.veq<1>
+    %1 = quake.alloca !quake.veq<1>
+    return
+  }
+
+// CHECK-LABEL:   func.func @test_veq_no_extract_refs3() {
+// CHECK:           %[[VAL_0:.*]] = quake.alloca !quake.veq<1>
+// CHECK:           %[[VAL_1:.*]] = quake.alloca !quake.veq<1>
+// CHECK:           %[[VAL_2:.*]] = quake.extract_ref %[[VAL_0]][0] : (!quake.veq<1>) -> !quake.ref
+// CHECK:           quake.h %[[VAL_2]] : (!quake.ref) -> ()
+// CHECK:           %[[VAL_3:.*]] = quake.extract_ref %[[VAL_1]][0] : (!quake.veq<1>) -> !quake.ref
+// CHECK:           quake.h %[[VAL_3]] : (!quake.ref) -> ()
+// CHECK:           %[[VAL_4:.*]] = quake.mz %[[VAL_2]] name "r00000" : (!quake.ref) -> !quake.measure
+// CHECK:           %[[VAL_5:.*]] = quake.mz %[[VAL_3]] name "r00001" : (!quake.ref) -> !quake.measure
+// CHECK:           return
+// CHECK:         }