From a1dda5adbec6a38e667ba0b1ef118ced283f71db Mon Sep 17 00:00:00 2001
From: Sachin Pisal <spisal@nvidia.com>
Date: Fri, 3 Jan 2025 13:12:59 -0800
Subject: [PATCH] moving the vector size check in get_state

Signed-off-by: Sachin Pisal <spisal@nvidia.com>
---
 .../examples/python/visualization.ipynb       | 24 +++++++++----------
 python/runtime/cudaq/algorithms/py_state.cpp  | 18 ++++++++++++++
 python/utils/OpaqueArguments.h                |  8 -------
 3 files changed, 30 insertions(+), 20 deletions(-)

diff --git a/docs/sphinx/examples/python/visualization.ipynb b/docs/sphinx/examples/python/visualization.ipynb
index 213d889f97..44a2f435dd 100644
--- 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
index e8d1ca2a98..a4c454a846 100644
--- 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
index 292751e73c..dd76e6fcc5 100644
--- 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);