moving the vector size check in get_state

Signed-off-by: Sachin Pisal <[email protected]>
NVIDIA · Jan 3, 2025 · a1dda5a · a1dda5a
1 parent 5281780
commit a1dda5a
Show file tree

Hide file tree

Showing 3 changed files with 30 additions and 20 deletions.
diff --git a/docs/sphinx/examples/python/visualization.ipynb b/docs/sphinx/examples/python/visualization.ipynb
@@ -35,7 +35,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 9,
+   "execution_count": 4,
    "metadata": {
     "scrolled": true
    },
@@ -64,7 +64,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 10,
+   "execution_count": 5,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -92,7 +92,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 11,
+   "execution_count": 6,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -113,7 +113,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 12,
+   "execution_count": 7,
    "metadata": {},
    "outputs": [
     {
@@ -140,7 +140,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 13,
+   "execution_count": 8,
    "metadata": {},
    "outputs": [
     {
@@ -167,7 +167,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 14,
+   "execution_count": 9,
    "metadata": {},
    "outputs": [
     {
@@ -194,7 +194,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 15,
+   "execution_count": 10,
    "metadata": {},
    "outputs": [
     {
@@ -221,7 +221,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 16,
+   "execution_count": 11,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -243,7 +243,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 17,
+   "execution_count": 12,
    "metadata": {},
    "outputs": [
     {
@@ -279,7 +279,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 18,
+   "execution_count": 13,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -302,7 +302,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 19,
+   "execution_count": 14,
    "metadata": {},
    "outputs": [
     {
@@ -328,7 +328,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 20,
+   "execution_count": 15,
    "metadata": {},
    "outputs": [
     {

diff --git a/python/runtime/cudaq/algorithms/py_state.cpp b/python/runtime/cudaq/algorithms/py_state.cpp
@@ -52,6 +52,24 @@ state pyGetState(py::object kernel, py::args args) {
 
   auto kernelName = kernel.attr("name").cast<std::string>();
   auto kernelMod = kernel.attr("module").cast<MlirModule>();
+
+  // Validate vector size if it is used for state initialization
+  if (args.size() == 1) {
+    py::object arg = args[0];
+
+    // Check if the argument is an iterable
+    if (py::hasattr(arg, "__len__")) {
+      size_t size = py::len(arg);
+
+      if (size == 0 || (size & (size - 1)) != 0) {
+        throw std::invalid_argument(
+            fmt::format("Invalid vector size: {}. When initializing a qubit "
+                        "register, the vector size must be a power of 2.",
+                        size));
+      }
+    }
+  }
+
   auto *argData = toOpaqueArgs(args, kernelMod, kernelName);
 
   return details::extractState([&]() mutable {

diff --git a/python/utils/OpaqueArguments.h b/python/utils/OpaqueArguments.h
@@ -99,14 +99,6 @@ inline py::args simplifiedValidateInputArguments(py::args &args) {
       if (shape.size() != 1)
         throw std::runtime_error("Cannot pass ndarray with shape != (N,).");
 
-      // Vector size validation
-      size_t size = shape[0].cast<size_t>();
-      if (size == 0 || (size & (size - 1)) != 0) {
-        throw std::invalid_argument(fmt::format(
-            "Invalid vector size: {}. Vector size must be a power of 2.",
-            size));
-      }
-
       arg = args[i].attr("tolist")();
     } else if (py::isinstance<py::str>(arg)) {
       arg = py::cast<std::string>(arg);