From c3ae059e931220351aadc9fa0e1f772b01c54f79 Mon Sep 17 00:00:00 2001
From: Eric Schweitz <eschweitz@nvidia.com>
Date: Mon, 18 Nov 2024 16:06:01 -0800
Subject: [PATCH] Fix issue #2382.

Adds a check to see if the execution context is observe to avoid the
bug mentioned in issue #2382. The issue mentions "bugs", so there may
be others.

Add the test case provided.

Signed-off-by: Eric Schweitz <eschweitz@nvidia.com>
---
 runtime/common/BaseRemoteRESTQPU.h |  4 ++-
 targettests/execution/observe.cpp  | 47 ++++++++++++++++++++++++++++++
 2 files changed, 50 insertions(+), 1 deletion(-)
 create mode 100644 targettests/execution/observe.cpp

diff --git a/runtime/common/BaseRemoteRESTQPU.h b/runtime/common/BaseRemoteRESTQPU.h
index 9f5af9d88b..062f27d46d 100644
--- a/runtime/common/BaseRemoteRESTQPU.h
+++ b/runtime/common/BaseRemoteRESTQPU.h
@@ -679,6 +679,8 @@ class BaseRemoteRESTQPU : public cudaq::QPU {
               }
             }
 
+            bool isObserve =
+                executionContext && executionContext->name == "observe";
             for (std::size_t i = 0; i < codes.size(); i++) {
               cudaq::ExecutionContext context("sample", localShots);
               context.reorderIdx = reorderIdx;
@@ -688,7 +690,7 @@ class BaseRemoteRESTQPU : public cudaq::QPU {
 
               // If there are multiple codes, this is likely a spin_op.
               // If so, use the code name instead of the global register.
-              if (codes.size() > 1) {
+              if (isObserve || (codes.size() > 1)) {
                 results.emplace_back(context.result.to_map(), codes[i].name);
                 results.back().sequentialData =
                     context.result.sequential_data();
diff --git a/targettests/execution/observe.cpp b/targettests/execution/observe.cpp
new file mode 100644
index 0000000000..2bdddd109d
--- /dev/null
+++ b/targettests/execution/observe.cpp
@@ -0,0 +1,47 @@
+/*******************************************************************************
+ * Copyright (c) 2022 - 2024 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.    *
+ ******************************************************************************/
+
+// RUN: nvq++ %cpp_std --target oqc --emulate %s -o %t && %t | FileCheck %s
+
+#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()(double theta) __qpu__ {
+    cudaq::qvector q(2);
+    x(q[0]);
+    ry(theta, q[1]);
+    x<cudaq::ctrl>(q[1], q[0]);
+  }
+};
+
+int main() {
+  // Build up your spin op algebraically
+  using namespace cudaq::spin;
+  cudaq::spin_op h = i(0) * z(1); // <-- MAKE THIS CHANGE HERE
+  // cudaq::spin_op h = 5.907 - 2.1433 * x(0) * x(1) - 2.1433 * y(0) * y(1) +
+  //                    .21829 * z(0) - 6.125 * z(1);
+
+  // 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
+  double energy = cudaq::observe(ansatz{}, h, .59);
+  printf("Attention.\n");
+  printf("Energy is %.16lf\n", energy);
+  printf("At ease.\n");
+  return 0;
+}
+
+// CHECK-LABEL: Attention
+// CHECK-NOT: Energy is 0.000000
+// CHECK-LABEL: At ease