feat: build autoqasm program directly from main decorator (#804)

examples/autoqasm/5_Pulse_programming_and_dynamical_decoupling.ipynb

@@ -37,7 +37,7 @@
     "\n",
     "import braket.experimental.autoqasm as aq\n",
     "from braket.experimental.autoqasm import pulse\n",
-    "from braket.experimental.autoqasm.instructions import rx, rz, measure\n",
+    "from braket.experimental.autoqasm.instructions import rx, rz\n",
     "\n",
     "from braket.aws import AwsDevice\n",
     "from braket.devices import Devices\n",
@@ -69,8 +69,11 @@
    "metadata": {},
    "outputs": [],
    "source": [
+    "qubit = 0\n",
+    "idle_duration = 10e-6\n",
+    "\n",
     "@aq.main\n",
-    "def idle(qubit: int, idle_duration: float):\n",
+    "def idle():\n",
     "    control_frame = device.frames[f\"q{qubit}_rf_frame\"]\n",
     "    pulse.delay(control_frame, idle_duration)\n",
     "    pulse.capture_v0(control_frame)"
@@ -104,7 +107,7 @@
     }
    ],
    "source": [
-    "print(idle(0, 10e-6).to_ir())"
+    "print(idle.to_ir())"
    ]
   },
   {
@@ -180,8 +183,12 @@
    "metadata": {},
    "outputs": [],
    "source": [
+    "qubit = 0\n",
+    "idle_duration = 10e-6\n",
+    "n_cycles = 3\n",
+    "\n",
     "@aq.main\n",
-    "def idle_with_dd(qubit: int, idle_duration: float, n_cycles: int = 1) -> None:\n",
+    "def idle_with_dd():\n",
     "    dd_spacing = idle_duration / (4*n_cycles)\n",
     "\n",
     "    control_frame = device.frames[f\"q{qubit}_rf_frame\"]\n",
@@ -247,7 +254,7 @@
     }
    ],
    "source": [
-    "print(idle_with_dd(0, 10e-6, 3).to_ir())"
+    "print(idle_with_dd.to_ir())"
    ]
   },
   {

examples/autoqasm/6_Customize_gate_calibrations.ipynb

@@ -79,7 +79,7 @@
      "text": [
       "OPENQASM 3.0;\n",
       "cal {\n",
-      "    waveform wf_drag_gaussian_1 = drag_gaussian(24.0ns, 2.547965400864ns, 2.3154335273287345e-10, 0.294418024251374, false);\n",
+      "    waveform wf_drag_gaussian_1 = drag_gaussian(24.0ns, 2.547965400864ns, 2.3223415460834803e-10, 0.28585537029609, false);\n",
       "    barrier $0;\n",
       "    shift_frequency(q0_rf_frame, -321047.14178613486);\n",
       "    play(q0_rf_frame, wf_drag_gaussian_1);\n",
@@ -158,12 +158,12 @@
      "text": [
       "Help on function rx in module braket.experimental.autoqasm.instructions.gates:\n",
       "\n",
-      "rx(target: Union[int, oqpy.classical_types._ClassicalVar, oqpy.base.OQPyExpression, str, oqpy.quantum_types.Qubit], angle: float) -> None\n",
+      "rx(target: Union[int, str, braket.registers.qubit.Qubit, oqpy.classical_types._ClassicalVar, oqpy.base.OQPyExpression, oqpy.quantum_types.Qubit], angle: Union[float, braket.parametric.free_parameter_expression.FreeParameterExpression, oqpy.classical_types._ClassicalVar]) -> None\n",
       "    X-axis rotation gate.\n",
       "    \n",
       "    Args:\n",
       "        target (QubitIdentifierType): Target qubit.\n",
-      "        angle (float): Rotation angle in radians.\n",
+      "        angle (GateParameterType): Rotation angle in radians.\n",
       "\n"
      ]
     }
@@ -234,8 +234,7 @@
     "    rz(\"$1\", 0.123)\n",
     "    measure(\"$0\")\n",
     "\n",
-    "main_program = my_program()\n",
-    "print(main_program.to_ir())"
+    "print(my_program.to_ir())"
    ]
   },
   {
@@ -277,7 +276,7 @@
     }
    ],
    "source": [
-    "custom_program = main_program.with_calibrations(my_rx_cal)\n",
+    "custom_program = my_program.with_calibrations(my_rx_cal)\n",
     "print(custom_program.to_ir())"
    ]
   },
@@ -347,7 +346,7 @@
     }
    ],
    "source": [
-    "custom_program = main_program.with_calibrations([my_rx_cal, my_rz_cal])\n",
+    "custom_program = my_program.with_calibrations([my_rx_cal, my_rz_cal])\n",
     "print(custom_program.to_ir())"
    ]
   },

src/braket/experimental/autoqasm/__init__.py

@@ -39,7 +39,7 @@ def my_program():
     result[0] = measure __qubits__[0];
     result[1] = measure __qubits__[1];
 """
-
+from . import errors, operators  # noqa: F401
 from .api import gate, gate_calibration, main, subroutine  # noqa: F401
 from .instructions import QubitIdentifierType as Qubit  # noqa: F401
 from .program import Program, build_program, verbatim  # noqa: F401

src/braket/experimental/autoqasm/api.py

@@ -42,16 +42,16 @@
 from braket.experimental.autoqasm.instructions.qubits import QubitIdentifierType as Qubit
 from braket.experimental.autoqasm.instructions.qubits import is_qubit_identifier_type
 from braket.experimental.autoqasm.program.gate_calibrations import GateCalibration
+from braket.parametric import FreeParameter
 
 
 def main(
     func: Optional[Callable] = None,
     *,
     num_qubits: Optional[int] = None,
     device: Optional[Union[Device, str]] = None,
-) -> Callable[..., aq_program.Program]:
-    """Decorator that converts a function into a callable that returns
-    a Program object containing the quantum program.
+) -> aq_program.Program | functools.partial:
+    """Decorator that converts a function into a Program object containing the quantum program.
 
     The decorator re-converts the target function whenever the decorated
     function is called, and a new Program object is returned each time.
@@ -65,13 +65,22 @@ def main(
             program. Can be either an Device object or a valid Amazon Braket device ARN.
 
     Returns:
-        Callable[..., Program]: A callable which returns the converted
-        quantum program when called.
+        Program | partial: The Program object containing the converted quantum program, or a
+        partial function of the `main` decorator.
     """
     if isinstance(device, str):
         device = AwsDevice(device)
 
-    return _function_wrapper(
+    # decorator is called on a Program
+    if isinstance(func, aq_program.Program):
+        return func
+
+    # decorator is used with parentheses
+    # (see _function_wrapper for more details)
+    if not (func and callable(func)):
+        return functools.partial(main, num_qubits=num_qubits, device=device)
+
+    program_builder = _function_wrapper(
         func,
         converter_callback=_convert_main,
         converter_args={
@@ -82,6 +91,8 @@ def main(
         },
     )
 
+    return program_builder()
+
 
 def subroutine(func: Optional[Callable] = None) -> Callable[..., aq_program.Program]:
     """Decorator that converts a function into a callable that will insert a subroutine into
@@ -182,7 +193,7 @@ def _wrapper(*args, **kwargs) -> Callable:
             optional_features=_autograph_optional_features(),
         )
         # Call the appropriate function converter
-        return converter_callback(func, options, args, kwargs, **converter_args)
+        return converter_callback(func, options=options, args=args, kwargs=kwargs, **converter_args)
 
     if inspect.isfunction(func) or inspect.ismethod(func):
         _wrapper = functools.update_wrapper(_wrapper, func)
@@ -204,7 +215,7 @@ def _convert_main(
     args: tuple[Any],
     kwargs: dict[str, Any],
     user_config: aq_program.UserConfig,
-) -> None:
+) -> aq_program.Program:
     """Convert the initial callable `f` into a full AutoQASM program `program`.
     Puts the contents of `f` at the global level of the program, rather than
     putting it into a subroutine as done in `_convert_subroutine`.
@@ -218,16 +229,25 @@ def _convert_main(
         args (tuple[Any]): Arguments passed to the program when called.
         kwargs (dict[str, Any]): Keyword arguments passed to the program when called.
         user_config (UserConfig): User-specified settings that influence program building.
+
+    Returns:
+        aq_program.Program: Generated AutoQASM Program.
     """
-    if aq_program.in_active_program_conversion_context():
-        raise errors.AutoQasmTypeError(
-            f"Cannot call main function '{f.__name__}' from another main function. Did you mean "
-            "to use '@aq.subroutine'?"
-        )
+    kwargs = {}
+    parameters = inspect.signature(f).parameters
 
     with aq_program.build_program(user_config) as program_conversion_context:
+        # Capture inputs to decorated function as `FreeParameter` inputs for the Program
+        for param in parameters.values():
+            if param.kind == param.POSITIONAL_OR_KEYWORD:
+                kwargs[param.name] = FreeParameter(param.name)
+                param_type = param.annotation if param.annotation is not param.empty else float
+                program_conversion_context.register_parameter(param.name, param_type)
+            else:
+                raise NotImplementedError
+
         # Process the program
-        aq_transpiler.converted_call(f, args, kwargs, options=options)
+        aq_transpiler.converted_call(f, (), kwargs, options=options)
 
         # Modify program to add global declarations if necessary
         _add_qubit_declaration(program_conversion_context)