Return correctly-sized bit register

amazon-braket · Sep 18, 2023 · 69175fd · 69175fd
1 parent b1418b6
commit 69175fd
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 6 deletions.
diff --git a/src/braket/experimental/autoqasm/api.py b/src/braket/experimental/autoqasm/api.py
@@ -437,10 +437,18 @@ def _make_return_instance_from_oqpy_return_type(return_type: Any) -> Any:
     if not return_type:
         return None
 
-    return_type = aq_types.conversions.var_type_from_ast_type(return_type)
-    if return_type == aq_types.ArrayVar:
+    var_type = aq_types.conversions.var_type_from_ast_type(return_type)
+    if var_type == aq_types.ArrayVar:
         return []
-    return return_type()
+    if var_type == aq_types.BitVar:
+        return var_type(size=_get_bitvar_size(return_type))
+    return var_type()
+
+
+def _get_bitvar_size(node: qasm_ast.BitType) -> Optional[int]:
+    if not isinstance(node, qasm_ast.BitType) or not node.size:
+        return None
+    return node.size.value
 
 
 def _convert_gate(

diff --git a/test/unit_tests/braket/experimental/autoqasm/test_api.py b/test/unit_tests/braket/experimental/autoqasm/test_api.py
@@ -483,9 +483,7 @@ def ground_state_measurements_subroutine() -> bit[3] {
     return __bit_0__;
 }
 qubit[6] __qubits__;
-"""
-    # TODO: this should be `bit[3]`, but there's a bug. It's being tracked in an issue.
-    expected += """bit __bit_1__;
+bit[3] __bit_1__ = "000";
 __bit_1__ = ground_state_measurements_subroutine();"""
     assert ground_state_measurements_wrapper().to_ir() == expected