Fix multi-dimensional ArrayVar declaration (#25)

amazon-braket · Jun 5, 2024 · 62ab4c5 · 62ab4c5
1 parent 90c2bbd
commit 62ab4c5
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 4 deletions.
diff --git a/src/autoqasm/types/types.py b/src/autoqasm/types/types.py
@@ -18,6 +18,7 @@
 from collections.abc import Iterable
 from typing import Any, List, Union, get_args
 
+import numpy as np
 import oqpy
 import oqpy.base
 from braket.circuits import FreeParameterExpression
@@ -114,7 +115,7 @@ def __init__(
         if not isinstance(init_expression, Iterable):
             raise errors.InvalidArrayDeclaration("init_expression must be an iterable type.")
 
-        dimensions = [len(init_expression)]
+        dimensions = np.shape(init_expression)
         super(ArrayVar, self).__init__(
             init_expression=init_expression,
             *args,

diff --git a/test/unit_tests/autoqasm/test_types.py b/test/unit_tests/autoqasm/test_types.py
@@ -756,9 +756,9 @@ def declare_array():
 def test_array_supports_multidimensional_arrays():
     @aq.main
     def declare_array():
-        aq.ArrayVar([[1, 2], [3, 4]])
+        a = aq.ArrayVar([[1, 2, 3], [4, 5, 6]])  # noqa: F841
 
     expected = """OPENQASM 3.0;
-array[int[32], 2, 2] a = {{1, 2}, {3, 4}};"""
+array[int[32], 2, 3] a = {{1, 2, 3}, {4, 5, 6}};"""
 
-    declare_array.build().to_ir() == expected
+    assert declare_array.build().to_ir() == expected