From 2a8bfcacb496530a220a38d0c7fa72cf21a3926d Mon Sep 17 00:00:00 2001
From: Lauren Capelluto <laurecap@amazon.com>
Date: Tue, 24 Oct 2023 15:04:12 -0400
Subject: [PATCH 01/10] feature: Add support for FreeParameters

---
 src/braket/devices/local_simulator.py         |   8 +-
 src/braket/experimental/autoqasm/api.py       |  13 +-
 .../autoqasm/instructions/gates.py            |  74 ++++-----
 .../autoqasm/instructions/instructions.py     |   6 +
 .../experimental/autoqasm/program/program.py  |  12 ++
 .../experimental/autoqasm/test_parameters.py  | 145 ++++++++++++++++++
 6 files changed, 220 insertions(+), 38 deletions(-)
 create mode 100644 test/unit_tests/braket/experimental/autoqasm/test_parameters.py

diff --git a/src/braket/devices/local_simulator.py b/src/braket/devices/local_simulator.py
index 253732a06..c3c70cd6d 100644
--- a/src/braket/devices/local_simulator.py
+++ b/src/braket/devices/local_simulator.py
@@ -317,7 +317,13 @@ def _(
         if DeviceActionType.OPENQASM not in simulator.properties.action:
             raise NotImplementedError(f"{type(simulator)} does not support OpenQASM programs")
         program = Program(source=program.to_ir(ir_type=IRType.OPENQASM))
-
+        if inputs:
+            inputs_copy = program.inputs.copy() if program.inputs is not None else {}
+            inputs_copy.update(inputs)
+            program = Program(
+                source=program.source,
+                inputs=inputs_copy,
+            )
         # Pass mcm=True to the simulator to enable mid-circuit measurement simulation.
         # When setting mcm=True, the simulator returns only the measurement results,
         # which we then wrap into a GateModelQuantumTaskResult object to return.
diff --git a/src/braket/experimental/autoqasm/api.py b/src/braket/experimental/autoqasm/api.py
index b6f421e3a..4f9d8d4c6 100644
--- a/src/braket/experimental/autoqasm/api.py
+++ b/src/braket/experimental/autoqasm/api.py
@@ -206,12 +206,23 @@ def _convert_main(
         # Process the program
         aq_transpiler.converted_call(f, args, kwargs, options=options)
 
-        # Modify program to add qubit declaration if necessary
+        # Modify program to add global declarations if necessary
         _add_qubit_declaration(program_conversion_context)
+        _add_io_declarations(program_conversion_context)
 
     return program_conversion_context.make_program()
 
 
+def _add_io_declarations(program_conversion_context: aq_program.ProgramConversionContext) -> None:
+    root_oqpy_program = program_conversion_context.get_oqpy_program(
+        scope=aq_program.ProgramScope.MAIN
+    )
+    root_oqpy_program.declare(
+        program_conversion_context.get_free_parameters(),
+        to_beginning=True,
+    )
+
+
 def _add_qubit_declaration(program_conversion_context: aq_program.ProgramConversionContext) -> None:
     """Modify the program to include a global qubit register declaration.
 
diff --git a/src/braket/experimental/autoqasm/instructions/gates.py b/src/braket/experimental/autoqasm/instructions/gates.py
index aff44f6b2..13b5ae13d 100644
--- a/src/braket/experimental/autoqasm/instructions/gates.py
+++ b/src/braket/experimental/autoqasm/instructions/gates.py
@@ -15,6 +15,8 @@
 """Quantum gates, unitary instructions, that apply to qubits.
 """
 
+from braket.circuits.free_parameter_expression import FreeParameterExpression
+
 from .instructions import _qubit_instruction
 from .qubits import QubitIdentifierType
 
@@ -52,14 +54,14 @@ def cnot(
 def cphaseshift(
     control: QubitIdentifierType,
     target: QubitIdentifierType,
-    angle: float,
+    angle: float | FreeParameterExpression,
 ) -> None:
     """Controlled phase shift gate.
 
     Args:
         control (QubitIdentifierType): Control qubit.
         target (QubitIdentifierType): Target qubit.
-        angle (float): Rotation angle in radians.
+        angle (float | FreeParameterExpression): Rotation angle in radians.
 
     """
     _qubit_instruction("cphaseshift", [control, target], angle)
@@ -68,14 +70,14 @@ def cphaseshift(
 def cphaseshift00(
     control: QubitIdentifierType,
     target: QubitIdentifierType,
-    angle: float,
+    angle: float | FreeParameterExpression,
 ) -> None:
     """Controlled phase shift gate for phasing the \\|00> state.
 
     Args:
         control (QubitIdentifierType): Control qubit.
         target (QubitIdentifierType): Target qubit.
-        angle (float): Rotation angle in radians.
+        angle (float | FreeParameterExpression): Rotation angle in radians.
 
     """
     _qubit_instruction("cphaseshift00", [control, target], angle)
@@ -84,14 +86,14 @@ def cphaseshift00(
 def cphaseshift01(
     control: QubitIdentifierType,
     target: QubitIdentifierType,
-    angle: float,
+    angle: float | FreeParameterExpression,
 ) -> None:
     """Controlled phase shift gate for phasing the \\|01> state.
 
     Args:
         control (QubitIdentifierType): Control qubit.
         target (QubitIdentifierType): Target qubit.
-        angle (float): Rotation angle in radians.
+        angle (float | FreeParameterExpression): Rotation angle in radians.
 
     """
     _qubit_instruction("cphaseshift01", [control, target], angle)
@@ -100,14 +102,14 @@ def cphaseshift01(
 def cphaseshift10(
     control: QubitIdentifierType,
     target: QubitIdentifierType,
-    angle: float,
+    angle: float | FreeParameterExpression,
 ) -> None:
     """Controlled phase shift gate for phasing the \\|10> state.
 
     Args:
         control (QubitIdentifierType): Control qubit.
         target (QubitIdentifierType): Target qubit.
-        angle (float): Rotation angle in radians.
+        angle (float | FreeParameterExpression): Rotation angle in radians.
 
     """
     _qubit_instruction("cphaseshift10", [control, target], angle)
@@ -187,13 +189,13 @@ def ecr(
 
 def gpi(
     target: QubitIdentifierType,
-    angle: float,
+    angle: float | FreeParameterExpression,
 ) -> None:
     """IonQ GPi gate.
 
     Args:
         target (QubitIdentifierType): Target qubit.
-        angle (float): Rotation angle in radians.
+        angle (float | FreeParameterExpression): Rotation angle in radians.
 
     """
     _qubit_instruction("gpi", [target], angle)
@@ -201,13 +203,13 @@ def gpi(
 
 def gpi2(
     target: QubitIdentifierType,
-    angle: float,
+    angle: float | FreeParameterExpression,
 ) -> None:
     """IonQ GPi2 gate.
 
     Args:
         target (QubitIdentifierType): Target qubit.
-        angle (float): Rotation angle in radians.
+        angle (float | FreeParameterExpression): Rotation angle in radians.
 
     """
     _qubit_instruction("gpi2", [target], angle)
@@ -254,18 +256,18 @@ def iswap(
 def ms(
     target_0: QubitIdentifierType,
     target_1: QubitIdentifierType,
-    angle_0: float,
-    angle_1: float,
-    angle_2: float,
+    angle_0: float | FreeParameterExpression,
+    angle_1: float | FreeParameterExpression,
+    angle_2: float | FreeParameterExpression,
 ) -> None:
     """IonQ Mølmer-Sørenson gate.
 
     Args:
         target_0 (QubitIdentifierType): Target qubit 0.
         target_1 (QubitIdentifierType): Target qubit 1.
-        angle_0 (float): Rotation angle 0 in radians.
-        angle_1 (float): Rotation angle 1 in radians.
-        angle_2 (float): Rotation angle 2 in radians.
+        angle_0 (float | FreeParameterExpression): Rotation angle 0 in radians.
+        angle_1 (float | FreeParameterExpression): Rotation angle 1 in radians.
+        angle_2 (float | FreeParameterExpression): Rotation angle 2 in radians.
 
     """
     _qubit_instruction("ms", [target_0, target_1], angle_0, angle_1, angle_2)
@@ -273,13 +275,13 @@ def ms(
 
 def phaseshift(
     target: QubitIdentifierType,
-    angle: float,
+    angle: float | FreeParameterExpression,
 ) -> None:
     """Phase shift gate.
 
     Args:
         target (QubitIdentifierType): Target qubit.
-        angle (float): Rotation angle in radians.
+        angle (float | FreeParameterExpression): Rotation angle in radians.
 
     """
     _qubit_instruction("phaseshift", [target], angle)
@@ -288,14 +290,14 @@ def phaseshift(
 def pswap(
     target_0: QubitIdentifierType,
     target_1: QubitIdentifierType,
-    angle: float,
+    angle: float | FreeParameterExpression,
 ) -> None:
     """PSwap gate.
 
     Args:
         target_0 (QubitIdentifierType): Target qubit 0.
         target_1 (QubitIdentifierType): Target qubit 1.
-        angle (float): Rotation angle in radians.
+        angle (float | FreeParameterExpression): Rotation angle in radians.
 
     """
     _qubit_instruction("pswap", [target_0, target_1], angle)
@@ -303,13 +305,13 @@ def pswap(
 
 def rx(
     target: QubitIdentifierType,
-    angle: float,
+    angle: float | FreeParameterExpression,
 ) -> None:
     """X-axis rotation gate.
 
     Args:
         target (QubitIdentifierType): Target qubit.
-        angle (float): Rotation angle in radians.
+        angle (float | FreeParameterExpression): Rotation angle in radians.
 
     """
     _qubit_instruction("rx", [target], angle)
@@ -317,13 +319,13 @@ def rx(
 
 def ry(
     target: QubitIdentifierType,
-    angle: float,
+    angle: float | FreeParameterExpression,
 ) -> None:
     """Y-axis rotation gate.
 
     Args:
         target (QubitIdentifierType): Target qubit.
-        angle (float): Rotation angle in radians.
+        angle (float | FreeParameterExpression): Rotation angle in radians.
 
     """
     _qubit_instruction("ry", [target], angle)
@@ -331,13 +333,13 @@ def ry(
 
 def rz(
     target: QubitIdentifierType,
-    angle: float,
+    angle: float | FreeParameterExpression,
 ) -> None:
     """Z-axis rotation gate.
 
     Args:
         target (QubitIdentifierType): Target qubit.
-        angle (float): Rotation angle in radians.
+        angle (float | FreeParameterExpression): Rotation angle in radians.
 
     """
     _qubit_instruction("rz", [target], angle)
@@ -444,14 +446,14 @@ def x(
 def xx(
     target_0: QubitIdentifierType,
     target_1: QubitIdentifierType,
-    angle: float,
+    angle: float | FreeParameterExpression,
 ) -> None:
     """Ising XX coupling gate.
 
     Args:
         target_0 (QubitIdentifierType): Target qubit 0.
         target_1 (QubitIdentifierType): Target qubit 1.
-        angle (float): Rotation angle in radians.
+        angle (float | FreeParameterExpression): Rotation angle in radians.
 
     """
     _qubit_instruction("xx", [target_0, target_1], angle)
@@ -460,14 +462,14 @@ def xx(
 def xy(
     target_0: QubitIdentifierType,
     target_1: QubitIdentifierType,
-    angle: float,
+    angle: float | FreeParameterExpression,
 ) -> None:
     """XY gates
 
     Args:
         target_0 (QubitIdentifierType): Target qubit 0.
         target_1 (QubitIdentifierType): Target qubit 1.
-        angle (float): Rotation angle in radians.
+        angle (float | FreeParameterExpression): Rotation angle in radians.
 
     """
     _qubit_instruction("xy", [target_0, target_1], angle)
@@ -488,14 +490,14 @@ def y(
 def yy(
     target_0: QubitIdentifierType,
     target_1: QubitIdentifierType,
-    angle: float,
+    angle: float | FreeParameterExpression,
 ) -> None:
     """Ising YY coupling gate.
 
     Args:
         target_0 (QubitIdentifierType): Target qubit 0.
         target_1 (QubitIdentifierType): Target qubit 1.
-        angle (float): Rotation angle in radians.
+        angle (float | FreeParameterExpression): Rotation angle in radians.
 
     """
     _qubit_instruction("yy", [target_0, target_1], angle)
@@ -516,14 +518,14 @@ def z(
 def zz(
     target_0: QubitIdentifierType,
     target_1: QubitIdentifierType,
-    angle: float,
+    angle: float | FreeParameterExpression,
 ) -> None:
     """Ising ZZ coupling gate.
 
     Args:
         target_0 (QubitIdentifierType): Target qubit 0.
         target_1 (QubitIdentifierType): Target qubit 1.
-        angle (float): Rotation angle in radians.
+        angle (float | FreeParameterExpression): Rotation angle in radians.
 
     """
     _qubit_instruction("zz", [target_0, target_1], angle)
diff --git a/src/braket/experimental/autoqasm/instructions/instructions.py b/src/braket/experimental/autoqasm/instructions/instructions.py
index 24077b16c..14d8db6e8 100644
--- a/src/braket/experimental/autoqasm/instructions/instructions.py
+++ b/src/braket/experimental/autoqasm/instructions/instructions.py
@@ -17,6 +17,7 @@
 
 from typing import Any
 
+from braket.circuits.free_parameter_expression import FreeParameterExpression
 from braket.experimental.autoqasm import program as aq_program
 
 from .qubits import QubitIdentifierType, _qubit
@@ -30,6 +31,11 @@ def _qubit_instruction(
 
     # Add the instruction to the program.
     program_conversion_context.register_gate(name)
+    for arg in args:
+        if isinstance(arg, FreeParameterExpression):
+            # TODO laurecap: Support for other types
+            # TODO laurecap: Support for expressions
+            program_conversion_context.register_parameter(arg.name)
     program_mode = aq_program.ProgramMode.UNITARY if is_unitary else aq_program.ProgramMode.NONE
     oqpy_program = program_conversion_context.get_oqpy_program(mode=program_mode)
     oqpy_program.gate([_qubit(q) for q in qubits], name, *args)
diff --git a/src/braket/experimental/autoqasm/program/program.py b/src/braket/experimental/autoqasm/program/program.py
index b6eaed22e..f3a507932 100644
--- a/src/braket/experimental/autoqasm/program/program.py
+++ b/src/braket/experimental/autoqasm/program/program.py
@@ -207,6 +207,7 @@ def __init__(self, user_config: Optional[UserConfig] = None):
         self._virtual_qubits_used = set()
         self._var_idx = 0
         self._has_pulse_control = False
+        self._free_parameters = {}
 
     def make_program(self) -> Program:
         """Makes a Program object using the oqpy program from this conversion context.
@@ -276,6 +277,17 @@ def register_gate(self, gate_name: str) -> None:
                     f"block. The native gates of the device are: {native_gates}"
                 )
 
+    def register_parameter(self, name: str) -> None:
+        """Register an input parameter with the given name, if it has not already been
+        registered. Only floats are currently supported.
+        """
+        if name not in self._free_parameters:
+            self._free_parameters[name] = oqpy.FloatVar("input", name=name)
+
+    def get_free_parameters(self) -> list[oqpy.FloatVar]:
+        """Return a list of named oqpy.Vars that are used as free parameters in the program."""
+        return list(self._free_parameters.values())
+
     def get_target_device(self) -> Optional[Device]:
         """Return the target device for the program, as specified by the user.
         Returns None if the user did not specify a target device.
diff --git a/test/unit_tests/braket/experimental/autoqasm/test_parameters.py b/test/unit_tests/braket/experimental/autoqasm/test_parameters.py
new file mode 100644
index 000000000..16970986b
--- /dev/null
+++ b/test/unit_tests/braket/experimental/autoqasm/test_parameters.py
@@ -0,0 +1,145 @@
+# Copyright Amazon.com Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You
+# may not use this file except in compliance with the License. A copy of
+# the License is located at
+#
+#     http://aws.amazon.com/apache2.0/
+#
+# or in the "license" file accompanying this file. This file is
+# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
+# ANY KIND, either express or implied. See the License for the specific
+# language governing permissions and limitations under the License.
+
+"""AutoQASM tests for FreeParameter support."""
+
+import braket.experimental.autoqasm as aq
+from braket.circuits import FreeParameter
+from braket.default_simulator import StateVectorSimulator
+from braket.devices.local_simulator import LocalSimulator
+from braket.experimental.autoqasm.instructions import cnot, h, measure, rx
+from braket.tasks.local_quantum_task import LocalQuantumTask
+
+
+def _test_on_local_sim(program: aq.Program) -> None:
+    device = LocalSimulator(backend=StateVectorSimulator())
+    task = device.run(program, shots=10)
+    assert isinstance(task, LocalQuantumTask)
+    assert isinstance(task.result().measurements, dict)
+
+
+def test_simple_parametric():
+    """Test a program with a parameter can be serialized."""
+
+    @aq.main
+    def parametric():
+        rx(0, FreeParameter("theta"))
+        measure(0)
+
+    expected = """OPENQASM 3.0;
+input float[64] theta;
+qubit[1] __qubits__;
+rx(theta) __qubits__[0];
+bit __bit_0__;
+__bit_0__ = measure __qubits__[0];"""
+    assert parametric().to_ir() == expected
+
+
+def test_multiple_parameters():
+    """TODO"""
+
+    @aq.main
+    def parametric():
+        rx(0, FreeParameter("alpha"))
+        rx(1, FreeParameter("theta"))
+        c = measure([0, 1])
+
+    expected = """OPENQASM 3.0;
+bit[2] c;
+input float[64] alpha;
+input float[64] theta;
+qubit[2] __qubits__;
+rx(alpha) __qubits__[0];
+rx(theta) __qubits__[1];
+bit[2] __bit_0__ = "00";
+__bit_0__[0] = measure __qubits__[0];
+__bit_0__[1] = measure __qubits__[1];
+c = __bit_0__;"""
+    assert parametric().to_ir() == expected
+
+
+def test_repeat_parameter():
+    """TODO"""
+
+    @aq.main
+    def parametric():
+        alpha = FreeParameter("alpha")
+        theta = FreeParameter("theta")
+        rx(0, alpha)
+        rx(1, theta)
+        cnot(0, 1)
+        rx(0, theta)
+        rx(1, alpha)
+
+    expected = """OPENQASM 3.0;
+input float[64] alpha;
+input float[64] theta;
+qubit[2] __qubits__;
+rx(alpha) __qubits__[0];
+rx(theta) __qubits__[1];
+cnot __qubits__[0], __qubits__[1];
+rx(theta) __qubits__[0];
+rx(alpha) __qubits__[1];"""
+    assert parametric().to_ir() == expected
+
+
+def test_parameter_in_subroutine():
+    """Test that parameters in subroutines are declared appropriately."""
+    pass
+
+
+def test_captured_parameter():
+    """Test that a parameter declared in a larger scope is captured
+    and functions correctly.
+    """
+    pass
+
+
+def test_parameter_expressions():
+    pass
+
+
+def test_parameter_in_args():
+    """TODO"""
+
+    @aq.main
+    def parametric(theta: float):
+        rx(0, theta)
+
+    parametric(FreeParameter("theta"))
+    # TODO
+
+
+def test_parameters_passed_as_args():
+    """Test that parameters work when passed as input values."""
+
+
+def test_multi_angle_gates():
+    """Test that FreeParameters work with gates that take multiple inputs."""
+    pass
+
+
+# TODO
+# - test parameter declared in subroutine input arg
+# - tests on local simulator
+# - other input types? local sim doesn't seem to support angle input types
+    # - integer inputs:
+        # input int basis; // 0 = X basis, 1 = Y basis, 2 = Z basis
+        # output bit result;
+        # qubit q;
+
+        # // Some complicated circuit...
+
+        # if (basis == 0) h q;
+        # else if (basis == 1) rx(π/2) q;
+        # result = measure q;

From 047451c3b62c4803e65aa7a1ece1e2488c004e25 Mon Sep 17 00:00:00 2001
From: Lauren Capelluto <laurecap@amazon.com>
Date: Wed, 25 Oct 2023 21:09:39 -0400
Subject: [PATCH 02/10] Add new tests and fix broken tests

Fixup lint
---
 src/braket/experimental/autoqasm/api.py       |  11 +-
 .../autoqasm/instructions/instructions.py     |   2 +-
 .../experimental/autoqasm/program/program.py  |  18 +-
 .../experimental/autoqasm/test_parameters.py  | 168 ++++++++++++++----
 4 files changed, 157 insertions(+), 42 deletions(-)

diff --git a/src/braket/experimental/autoqasm/api.py b/src/braket/experimental/autoqasm/api.py
index 4f9d8d4c6..5b56eb992 100644
--- a/src/braket/experimental/autoqasm/api.py
+++ b/src/braket/experimental/autoqasm/api.py
@@ -18,7 +18,7 @@
 import inspect
 from collections.abc import Callable
 from types import FunctionType
-from typing import Any, Optional, Union
+from typing import Any, Optional, Union, get_args
 
 import openqasm3.ast as qasm_ast
 import oqpy.base
@@ -545,9 +545,12 @@ def _get_gate_args(f: Callable) -> aq_program.GateArgs:
             )
 
         if param.annotation == aq_instructions.QubitIdentifierType:
-            gate_args.append(param.name, True)
-        elif param.annotation in [float, aq_types.FloatVar]:
-            gate_args.append(param.name, False)
+            gate_args.append_qubit(param.name)
+        elif param.annotation in [float, aq_types.FloatVar] or (
+            get_args(param.annotation)
+            and any(type_ in [float, aq_types.FloatVar] for type_ in get_args(param.annotation))
+        ):
+            gate_args.append_angle(param.name)
         else:
             raise errors.ParameterTypeError(
                 f'Parameter "{param.name}" for gate "{f.__name__}" '
diff --git a/src/braket/experimental/autoqasm/instructions/instructions.py b/src/braket/experimental/autoqasm/instructions/instructions.py
index 14d8db6e8..f5e1a9b05 100644
--- a/src/braket/experimental/autoqasm/instructions/instructions.py
+++ b/src/braket/experimental/autoqasm/instructions/instructions.py
@@ -33,7 +33,7 @@ def _qubit_instruction(
     program_conversion_context.register_gate(name)
     for arg in args:
         if isinstance(arg, FreeParameterExpression):
-            # TODO laurecap: Support for other types
+            # TODO laurecap: Support for integers
             # TODO laurecap: Support for expressions
             program_conversion_context.register_parameter(arg.name)
     program_mode = aq_program.ProgramMode.UNITARY if is_unitary else aq_program.ProgramMode.NONE
diff --git a/src/braket/experimental/autoqasm/program/program.py b/src/braket/experimental/autoqasm/program/program.py
index f3a507932..254708d9c 100644
--- a/src/braket/experimental/autoqasm/program/program.py
+++ b/src/braket/experimental/autoqasm/program/program.py
@@ -162,17 +162,21 @@ def __init__(self):
     def __len__(self):
         return len(self._args)
 
-    def append(self, name: str, is_qubit: bool) -> None:
-        """Appends an argument to the list of gate arguments.
+    def append_qubit(self, name: str) -> None:
+        """Appends a qubit argument to the list of gate arguments.
 
         Args:
             name (str): The name of the argument.
-            is_qubit (bool): Whether the argument represents a qubit.
         """
-        if is_qubit:
-            self._args.append(oqpy.Qubit(name, needs_declaration=False))
-        else:
-            self._args.append(oqpy.AngleVar(name=name))
+        self._args.append(oqpy.Qubit(name, needs_declaration=False))
+
+    def append_angle(self, name: str) -> None:
+        """Appends a parameter argument to the list of gate arguments.
+
+        Args:
+            name (str): The name of the argument.
+        """
+        self._args.append(oqpy.AngleVar(name=name))
 
     @property
     def qubits(self) -> list[oqpy.Qubit]:
diff --git a/test/unit_tests/braket/experimental/autoqasm/test_parameters.py b/test/unit_tests/braket/experimental/autoqasm/test_parameters.py
index 16970986b..16602f2ea 100644
--- a/test/unit_tests/braket/experimental/autoqasm/test_parameters.py
+++ b/test/unit_tests/braket/experimental/autoqasm/test_parameters.py
@@ -17,7 +17,7 @@
 from braket.circuits import FreeParameter
 from braket.default_simulator import StateVectorSimulator
 from braket.devices.local_simulator import LocalSimulator
-from braket.experimental.autoqasm.instructions import cnot, h, measure, rx
+from braket.experimental.autoqasm.instructions import cnot, cphaseshift, gpi, h, measure, ms, rx, rz
 from braket.tasks.local_quantum_task import LocalQuantumTask
 
 
@@ -46,13 +46,13 @@ def parametric():
 
 
 def test_multiple_parameters():
-    """TODO"""
+    """Test that multiple free parameters all appear in the processed program."""
 
     @aq.main
     def parametric():
         rx(0, FreeParameter("alpha"))
         rx(1, FreeParameter("theta"))
-        c = measure([0, 1])
+        c = measure([0, 1])  # noqa: F841
 
     expected = """OPENQASM 3.0;
 bit[2] c;
@@ -69,7 +69,7 @@ def parametric():
 
 
 def test_repeat_parameter():
-    """TODO"""
+    """Test that programs can use the same parameter multiple times."""
 
     @aq.main
     def parametric():
@@ -95,51 +95,159 @@ def parametric():
 
 def test_parameter_in_subroutine():
     """Test that parameters in subroutines are declared appropriately."""
-    pass
+
+    @aq.subroutine
+    def rx_alpha(qubit: int):
+        rx(qubit, FreeParameter("alpha"))
+
+    @aq.main(num_qubits=3)
+    def parametric():
+        rx_alpha(2)
+
+    expected = """OPENQASM 3.0;
+def rx_alpha(int[32] qubit) {
+    rx(alpha) __qubits__[qubit];
+}
+input float[64] alpha;
+qubit[3] __qubits__;
+rx_alpha(2);"""
+    assert parametric().to_ir() == expected
 
 
 def test_captured_parameter():
     """Test that a parameter declared in a larger scope is captured
     and functions correctly.
     """
-    pass
+
+    alpha = FreeParameter("alpha")
+
+    @aq.main
+    def parametric():
+        rz(0, alpha)
+        rx(1, alpha)
+
+    expected = """OPENQASM 3.0;
+input float[64] alpha;
+qubit[2] __qubits__;
+rz(alpha) __qubits__[0];
+rx(alpha) __qubits__[1];"""
+    assert parametric().to_ir() == expected
+
+
+def test_multi_angle_gates():
+    """Test that FreeParameters work with gates that take multiple inputs."""
+
+    @aq.main(num_qubits=5)
+    def parametric(qubit_0: int, phi: float, theta: float):
+        ms(0, qubit_0, phi, phi, theta)
+
+    expected = """OPENQASM 3.0;
+input float[64] phi;
+qubit[5] __qubits__;
+ms(phi, phi, 0.5) __qubits__[0], __qubits__[2];"""
+    assert parametric(2, FreeParameter("phi"), 0.5).to_ir() == expected
+
+
+def test_parameters_passed_as_main_arg():
+    """Test that parameters work when passed as input values."""
+
+    @aq.main
+    def parametric(phi: float):
+        cphaseshift(0, 1, phi)
+
+    expected = """OPENQASM 3.0;
+input float[64] my_phi;
+qubit[2] __qubits__;
+cphaseshift(my_phi) __qubits__[0], __qubits__[1];"""
+    assert parametric(FreeParameter("my_phi")).to_ir() == expected
+
+
+def test_parameters_passed_as_subroutine_arg():
+    """Test that parameters work when passed as input values."""
+    # FIXME
+
+    @aq.subroutine
+    def silly_ms(qubit_0: int, phi: float, theta: float):
+        ms(0, qubit_0, phi, phi, theta)
+
+    @aq.main(num_qubits=5)
+    def parametric():
+        silly_ms(1, FreeParameter("alpha"), 0.707)
+        silly_ms(3, 0.5, FreeParameter("beta"))
+
+    expected = """OPENQASM 3.0;
+def silly_ms(int[32] qubit_0, float[64] phi, float[64] theta) {
+    ms(phi, phi, theta) __qubits__[0], __qubits__[qubit_0];
+}
+input float[64] alpha;
+input float[64] beta;
+qubit[5] __qubits__;
+silly_ms(1, alpha, 0.707);
+silly_ms(3, 0.5, beta);"""
+    assert parametric().to_ir() == expected
 
 
 def test_parameter_expressions():
-    pass
+    """Test expressions of free parameters with numeric literals."""
 
+    @aq.main
+    def parametric():
+        expr = (2 * FreeParameter("theta")) + 1.5
+        rx(0, expr)
 
-def test_parameter_in_args():
-    """TODO"""
+    # TODO
+    expected = """OPENQASM 3.0;"""
+    assert parametric().to_ir() == expected
+
+
+def test_multi_parameter_expressions():
+    """Test expresssions of multiple free parameters."""
 
     @aq.main
-    def parametric(theta: float):
-        rx(0, theta)
+    def parametric():
+        expr = FreeParameter("alpha") * FreeParameter("theta")
+        gpi(0, expr)
 
-    parametric(FreeParameter("theta"))
     # TODO
+    expected = """OPENQASM 3.0;"""
+    assert parametric().to_ir() == expected
 
 
-def test_parameters_passed_as_args():
-    """Test that parameters work when passed as input values."""
+def test_integer_parameters():
+    """Test integer input parameter type."""
 
+    @aq.main
+    def parametric(qubit: int):
+        basis = FreeParameter("basis")
+        if basis == 0:
+            h(qubit)
+        elif basis == 1:
+            rx(qubit, 0.5)
+        else:
+            pass
+        result = measure(qubit)  # noqa: F841
 
-def test_multi_angle_gates():
-    """Test that FreeParameters work with gates that take multiple inputs."""
-    pass
+    # TODO
+    expected = """OPENQASM 3.0;
+bit result;
+input int[32] basis;
+qubit[2] __qubits__;
+if (basis == 0) {
+    h __qubits__[1];
+} else if (basis == 1) {
+    rx(0.5) __qubits__[1];
+}
+bit __bit_0__;
+__bit_0__ = measure __qubits__[1];
+result = __bit_0__;"""
+    assert parametric(1).to_ir() == expected
+
+
+def test_execution():
+    """TODO"""
+    # tests on local simulator
 
 
 # TODO
-# - test parameter declared in subroutine input arg
-# - tests on local simulator
-# - other input types? local sim doesn't seem to support angle input types
-    # - integer inputs:
-        # input int basis; // 0 = X basis, 1 = Y basis, 2 = Z basis
-        # output bit result;
-        # qubit q;
-
-        # // Some complicated circuit...
-
-        # if (basis == 0) h q;
-        # else if (basis == 1) rx(π/2) q;
-        # result = measure q;
+# - local sim doesn't seem to support angle input types
+# TODO: gate args?

From 3896cc477b449ec028be1e97e6e497afac428692 Mon Sep 17 00:00:00 2001
From: Lauren Capelluto <laurecap@amazon.com>
Date: Thu, 26 Oct 2023 10:37:35 -0400
Subject: [PATCH 03/10] Add more tests

---
 .../autoqasm/instructions/instructions.py     |  11 +-
 .../experimental/autoqasm/program/program.py  |   3 +
 .../experimental/autoqasm/test_parameters.py  | 139 +++++++++++-------
 3 files changed, 97 insertions(+), 56 deletions(-)

diff --git a/src/braket/experimental/autoqasm/instructions/instructions.py b/src/braket/experimental/autoqasm/instructions/instructions.py
index f5e1a9b05..3e8c4c805 100644
--- a/src/braket/experimental/autoqasm/instructions/instructions.py
+++ b/src/braket/experimental/autoqasm/instructions/instructions.py
@@ -12,11 +12,11 @@
 # language governing permissions and limitations under the License.
 
 
-"""Non-unitary instructions that apply to qubits.
-"""
+"""Non-unitary instructions that apply to qubits."""
 
 from typing import Any
 
+from braket.circuits.free_parameter import FreeParameter
 from braket.circuits.free_parameter_expression import FreeParameterExpression
 from braket.experimental.autoqasm import program as aq_program
 
@@ -32,10 +32,11 @@ def _qubit_instruction(
     # Add the instruction to the program.
     program_conversion_context.register_gate(name)
     for arg in args:
-        if isinstance(arg, FreeParameterExpression):
-            # TODO laurecap: Support for integers
-            # TODO laurecap: Support for expressions
+        if isinstance(arg, FreeParameter):
             program_conversion_context.register_parameter(arg.name)
+        elif isinstance(arg, FreeParameterExpression):
+            # TODO laurecap: Support for expressions
+            raise NotImplementedError("Expressions of FreeParameters will be supported shortly!")
     program_mode = aq_program.ProgramMode.UNITARY if is_unitary else aq_program.ProgramMode.NONE
     oqpy_program = program_conversion_context.get_oqpy_program(mode=program_mode)
     oqpy_program.gate([_qubit(q) for q in qubits], name, *args)
diff --git a/src/braket/experimental/autoqasm/program/program.py b/src/braket/experimental/autoqasm/program/program.py
index 254708d9c..c8265a10d 100644
--- a/src/braket/experimental/autoqasm/program/program.py
+++ b/src/braket/experimental/autoqasm/program/program.py
@@ -284,6 +284,9 @@ def register_gate(self, gate_name: str) -> None:
     def register_parameter(self, name: str) -> None:
         """Register an input parameter with the given name, if it has not already been
         registered. Only floats are currently supported.
+
+        Args:
+            name (str): The identifier for the parameter.
         """
         if name not in self._free_parameters:
             self._free_parameters[name] = oqpy.FloatVar("input", name=name)
diff --git a/test/unit_tests/braket/experimental/autoqasm/test_parameters.py b/test/unit_tests/braket/experimental/autoqasm/test_parameters.py
index 16602f2ea..6cf50d79d 100644
--- a/test/unit_tests/braket/experimental/autoqasm/test_parameters.py
+++ b/test/unit_tests/braket/experimental/autoqasm/test_parameters.py
@@ -11,7 +11,9 @@
 # ANY KIND, either express or implied. See the License for the specific
 # language governing permissions and limitations under the License.
 
-"""AutoQASM tests for FreeParameter support."""
+"""AutoQASM tests for parameter support."""
+
+import pytest
 
 import braket.experimental.autoqasm as aq
 from braket.circuits import FreeParameter
@@ -21,35 +23,68 @@
 from braket.tasks.local_quantum_task import LocalQuantumTask
 
 
-def _test_on_local_sim(program: aq.Program) -> None:
+def _test_parametric_on_local_sim(program: aq.Program, inputs: dict[str, float]) -> None:
     device = LocalSimulator(backend=StateVectorSimulator())
-    task = device.run(program, shots=10)
+    task = device.run(program, shots=10, inputs=inputs)
     assert isinstance(task, LocalQuantumTask)
     assert isinstance(task.result().measurements, dict)
+    return task.result().measurements
+
+
+@aq.main
+def simple_parametric():
+    rx(0, FreeParameter("theta"))
+    measure(0)
 
 
 def test_simple_parametric():
     """Test a program with a parameter can be serialized."""
 
-    @aq.main
-    def parametric():
-        rx(0, FreeParameter("theta"))
-        measure(0)
-
     expected = """OPENQASM 3.0;
 input float[64] theta;
 qubit[1] __qubits__;
 rx(theta) __qubits__[0];
 bit __bit_0__;
 __bit_0__ = measure __qubits__[0];"""
-    assert parametric().to_ir() == expected
+    assert simple_parametric().to_ir() == expected
+
+
+def test_sim_simple():
+    measurements = _test_parametric_on_local_sim(simple_parametric(), {"theta": 0})
+    assert 1 not in measurements["__bit_0__"]
+    measurements = _test_parametric_on_local_sim(simple_parametric(), {"theta": 3.14})
+    assert 0 not in measurements["__bit_0__"]
+
+
+@aq.main
+def multi_parametric():
+    rx(0, FreeParameter("alpha"))
+    rx(1, FreeParameter("theta"))
+    c = measure([0, 1])  # noqa: F841
 
 
 def test_multiple_parameters():
     """Test that multiple free parameters all appear in the processed program."""
 
+    expected = """OPENQASM 3.0;
+bit[2] c;
+input float[64] alpha;
+input float[64] theta;
+qubit[2] __qubits__;
+rx(alpha) __qubits__[0];
+rx(theta) __qubits__[1];
+bit[2] __bit_0__ = "00";
+__bit_0__[0] = measure __qubits__[0];
+__bit_0__[1] = measure __qubits__[1];
+c = __bit_0__;"""
+    assert multi_parametric().to_ir() == expected
+
+
+def test_typed_parameters():
+    """Test that multiple free parameters all appear in the processed program."""
+
     @aq.main
-    def parametric():
+    def multi_parametric():
         rx(0, FreeParameter("alpha"))
         rx(1, FreeParameter("theta"))
         c = measure([0, 1])  # noqa: F841
@@ -65,7 +100,14 @@ def parametric():
 __bit_0__[0] = measure __qubits__[0];
 __bit_0__[1] = measure __qubits__[1];
 c = __bit_0__;"""
-    assert parametric().to_ir() == expected
+    assert multi_parametric().to_ir() == expected
+
+
+def test_sim_multi_param():
+    measurements = _test_parametric_on_local_sim(multi_parametric(), {"alpha": 3.14, "theta": 0})
+    assert all(val == "10" for val in measurements["c"])
+    measurements = _test_parametric_on_local_sim(multi_parametric(), {"alpha": 0, "theta": 3.14})
+    assert all(val == "01" for val in measurements["c"])
 
 
 def test_repeat_parameter():
@@ -148,6 +190,14 @@ def parametric(qubit_0: int, phi: float, theta: float):
     assert parametric(2, FreeParameter("phi"), 0.5).to_ir() == expected
 
 
+def test_sim_multi_angle():
+    @aq.main
+    def parametric(phi: float, theta: float):
+        ms(0, 1, phi, phi, theta)
+
+    _test_parametric_on_local_sim(parametric(FreeParameter("phi"), 0.0), {"phi": 3.14})
+
+
 def test_parameters_passed_as_main_arg():
     """Test that parameters work when passed as input values."""
 
@@ -187,19 +237,43 @@ def silly_ms(int[32] qubit_0, float[64] phi, float[64] theta) {
     assert parametric().to_ir() == expected
 
 
+def test_sim_subroutine_arg():
+    @aq.subroutine
+    def rx_theta(theta: float):
+        rx(0, theta)
+
+    @aq.main
+    def parametric():
+        rx_theta(FreeParameter("theta"))
+
+    # TODO
+    measurements = _test_parametric_on_local_sim(parametric(), {"theta": 3.14})
+    assert 0 not in measurements["__bit_0__"]
+
+
 def test_parameter_expressions():
     """Test expressions of free parameters with numeric literals."""
 
     @aq.main
     def parametric():
-        expr = (2 * FreeParameter("theta")) + 1.5
-        rx(0, expr)
+        expr = 2 * FreeParameter("theta")
+        gpi(0, expr)
 
     # TODO
     expected = """OPENQASM 3.0;"""
     assert parametric().to_ir() == expected
 
 
+def test_sim_expressions():
+    @aq.main
+    def parametric():
+        rx(0, 2 * FreeParameter("phi"))
+        measure(0)
+
+    measurements = _test_parametric_on_local_sim(parametric(), {"phi": 3.14 / 2})
+    assert 0 not in measurements["__bit_0__"]
+
+
 def test_multi_parameter_expressions():
     """Test expresssions of multiple free parameters."""
 
@@ -213,41 +287,4 @@ def parametric():
     assert parametric().to_ir() == expected
 
 
-def test_integer_parameters():
-    """Test integer input parameter type."""
-
-    @aq.main
-    def parametric(qubit: int):
-        basis = FreeParameter("basis")
-        if basis == 0:
-            h(qubit)
-        elif basis == 1:
-            rx(qubit, 0.5)
-        else:
-            pass
-        result = measure(qubit)  # noqa: F841
-
-    # TODO
-    expected = """OPENQASM 3.0;
-bit result;
-input int[32] basis;
-qubit[2] __qubits__;
-if (basis == 0) {
-    h __qubits__[1];
-} else if (basis == 1) {
-    rx(0.5) __qubits__[1];
-}
-bit __bit_0__;
-__bit_0__ = measure __qubits__[1];
-result = __bit_0__;"""
-    assert parametric(1).to_ir() == expected
-
-
-def test_execution():
-    """TODO"""
-    # tests on local simulator
-
-
-# TODO
-# - local sim doesn't seem to support angle input types
-# TODO: gate args?
+# TODO: test gate args

From 1dd5754d50dff449694f114fe45de154247c0537 Mon Sep 17 00:00:00 2001
From: Lauren Capelluto <laurecap@amazon.com>
Date: Thu, 26 Oct 2023 17:48:37 -0400
Subject: [PATCH 04/10] Scope down to only floats

---
 .../experimental/autoqasm/test_parameters.py  | 73 +++++++++++--------
 1 file changed, 43 insertions(+), 30 deletions(-)

diff --git a/test/unit_tests/braket/experimental/autoqasm/test_parameters.py b/test/unit_tests/braket/experimental/autoqasm/test_parameters.py
index 6cf50d79d..8b2a07289 100644
--- a/test/unit_tests/braket/experimental/autoqasm/test_parameters.py
+++ b/test/unit_tests/braket/experimental/autoqasm/test_parameters.py
@@ -19,7 +19,17 @@
 from braket.circuits import FreeParameter
 from braket.default_simulator import StateVectorSimulator
 from braket.devices.local_simulator import LocalSimulator
-from braket.experimental.autoqasm.instructions import cnot, cphaseshift, gpi, h, measure, ms, rx, rz
+from braket.experimental.autoqasm.instructions import (
+    cnot,
+    cphaseshift,
+    gpi,
+    h,
+    measure,
+    ms,
+    rx,
+    rz,
+    x,
+)
 from braket.tasks.local_quantum_task import LocalQuantumTask
 
 
@@ -214,7 +224,6 @@ def parametric(phi: float):
 
 def test_parameters_passed_as_subroutine_arg():
     """Test that parameters work when passed as input values."""
-    # FIXME
 
     @aq.subroutine
     def silly_ms(qubit_0: int, phi: float, theta: float):
@@ -246,45 +255,49 @@ def rx_theta(theta: float):
     def parametric():
         rx_theta(FreeParameter("theta"))
 
-    # TODO
     measurements = _test_parametric_on_local_sim(parametric(), {"theta": 3.14})
     assert 0 not in measurements["__bit_0__"]
 
 
-def test_parameter_expressions():
-    """Test expressions of free parameters with numeric literals."""
+def test_parameter_as_condition():
+    """Test parameters used in conditional statements."""
 
     @aq.main
-    def parametric():
-        expr = 2 * FreeParameter("theta")
-        gpi(0, expr)
-
-    # TODO
-    expected = """OPENQASM 3.0;"""
-    assert parametric().to_ir() == expected
-
-
-def test_sim_expressions():
-    @aq.main
-    def parametric():
-        rx(0, 2 * FreeParameter("phi"))
+    def parametric(val: float):
+        threshold = 0.9
+        if val > threshold:
+            x(0)
         measure(0)
 
-    measurements = _test_parametric_on_local_sim(parametric(), {"phi": 3.14 / 2})
-    assert 0 not in measurements["__bit_0__"]
+    expected = """OPENQASM 3.0;
+input float[64] val;
+float[64] threshold;
+qubit[1] __qubits__;
+threshold = 0.9;
+if (val > threshold) {
+    x __qubits__[0];
+}
+bit __bit_0__;
+__bit_0__ = measure __qubits__[0];"""
+    assert parametric(FreeParameter("val")).to_ir() == expected
 
 
-def test_multi_parameter_expressions():
-    """Test expresssions of multiple free parameters."""
+def test_parametric_gate_args():
+    """Test that gates can be used with parameters."""
 
-    @aq.main
+    @aq.gate
+    def rx_theta(q: aq.Qubit, theta: float):
+        rx(q, theta)
+
+    @aq.main(num_qubits=3)
     def parametric():
-        expr = FreeParameter("alpha") * FreeParameter("theta")
-        gpi(0, expr)
+        rx_theta(2, FreeParameter("θ"))
 
-    # TODO
-    expected = """OPENQASM 3.0;"""
+    expected = """OPENQASM 3.0;
+gate rx_theta(theta) q {
+    rx(theta) q;
+}
+input float[64] θ;
+qubit[3] __qubits__;
+rx_theta(θ) __qubits__[2];"""
     assert parametric().to_ir() == expected
-
-
-# TODO: test gate args

From a77de15c1bdcf58994b35204d804fac75dd45a21 Mon Sep 17 00:00:00 2001
From: Lauren Capelluto <laurecap@amazon.com>
Date: Fri, 27 Oct 2023 12:33:08 -0400
Subject: [PATCH 05/10] Fix subroutine argument processing

---
 src/braket/experimental/autoqasm/api.py       |  1 +
 .../autoqasm/instructions/instructions.py     |  9 +-------
 .../experimental/autoqasm/program/program.py  | 17 ++++++++++++++
 .../experimental/autoqasm/test_parameters.py  | 22 +++++++++++++++++++
 4 files changed, 41 insertions(+), 8 deletions(-)

diff --git a/src/braket/experimental/autoqasm/api.py b/src/braket/experimental/autoqasm/api.py
index 5b56eb992..2acea01bb 100644
--- a/src/braket/experimental/autoqasm/api.py
+++ b/src/braket/experimental/autoqasm/api.py
@@ -307,6 +307,7 @@ def _convert_subroutine(
 
             # Process the program
             subroutine_function_call = oqpy_sub(oqpy_program, *args, **kwargs)
+            program_conversion_context.register_args(args)
 
             # Mark that we are finished processing this function
             program_conversion_context.subroutines_processing.remove(f)
diff --git a/src/braket/experimental/autoqasm/instructions/instructions.py b/src/braket/experimental/autoqasm/instructions/instructions.py
index 3e8c4c805..6a846caf0 100644
--- a/src/braket/experimental/autoqasm/instructions/instructions.py
+++ b/src/braket/experimental/autoqasm/instructions/instructions.py
@@ -16,8 +16,6 @@
 
 from typing import Any
 
-from braket.circuits.free_parameter import FreeParameter
-from braket.circuits.free_parameter_expression import FreeParameterExpression
 from braket.experimental.autoqasm import program as aq_program
 
 from .qubits import QubitIdentifierType, _qubit
@@ -31,12 +29,7 @@ def _qubit_instruction(
 
     # Add the instruction to the program.
     program_conversion_context.register_gate(name)
-    for arg in args:
-        if isinstance(arg, FreeParameter):
-            program_conversion_context.register_parameter(arg.name)
-        elif isinstance(arg, FreeParameterExpression):
-            # TODO laurecap: Support for expressions
-            raise NotImplementedError("Expressions of FreeParameters will be supported shortly!")
+    program_conversion_context.register_args(args)
     program_mode = aq_program.ProgramMode.UNITARY if is_unitary else aq_program.ProgramMode.NONE
     oqpy_program = program_conversion_context.get_oqpy_program(mode=program_mode)
     oqpy_program.gate([_qubit(q) for q in qubits], name, *args)
diff --git a/src/braket/experimental/autoqasm/program/program.py b/src/braket/experimental/autoqasm/program/program.py
index c8265a10d..d2538061c 100644
--- a/src/braket/experimental/autoqasm/program/program.py
+++ b/src/braket/experimental/autoqasm/program/program.py
@@ -23,6 +23,8 @@
 
 import oqpy.base
 
+from braket.circuits.free_parameter import FreeParameter
+from braket.circuits.free_parameter_expression import FreeParameterExpression
 from braket.circuits.serialization import IRType, SerializableProgram
 from braket.device_schema import DeviceActionType
 from braket.devices.device import Device
@@ -281,6 +283,21 @@ def register_gate(self, gate_name: str) -> None:
                     f"block. The native gates of the device are: {native_gates}"
                 )
 
+    def register_args(self, args: list[Any]) -> None:
+        """Register any FreeParameters in the list of arguments.
+
+        Args:
+            args (list[Any]): Arguments passed to the main program or a subroutine.
+        """
+        for arg in args:
+            if isinstance(arg, FreeParameter):
+                self.register_parameter(arg.name)
+            elif isinstance(arg, FreeParameterExpression):
+                # TODO laurecap: Support for expressions
+                raise NotImplementedError(
+                    "Expressions of FreeParameters will be supported shortly!"
+                )
+
     def register_parameter(self, name: str) -> None:
         """Register an input parameter with the given name, if it has not already been
         registered. Only floats are currently supported.
diff --git a/test/unit_tests/braket/experimental/autoqasm/test_parameters.py b/test/unit_tests/braket/experimental/autoqasm/test_parameters.py
index 8b2a07289..0d4f9dc90 100644
--- a/test/unit_tests/braket/experimental/autoqasm/test_parameters.py
+++ b/test/unit_tests/braket/experimental/autoqasm/test_parameters.py
@@ -222,6 +222,27 @@ def parametric(phi: float):
     assert parametric(FreeParameter("my_phi")).to_ir() == expected
 
 
+def test_simple_subroutine_arg():
+    """Test that parameters work when passed as input values."""
+
+    @aq.subroutine
+    def silly_rz(theta: float):
+        rz(0, theta)
+
+    @aq.main(num_qubits=1)
+    def parametric():
+        silly_rz(FreeParameter("alpha"))
+
+    expected = """OPENQASM 3.0;
+def silly_rz(float[64] theta) {
+    rz(theta) __qubits__[0];
+}
+input float[64] alpha;
+qubit[1] __qubits__;
+silly_rz(alpha);"""
+    assert parametric().to_ir() == expected
+
+
 def test_parameters_passed_as_subroutine_arg():
     """Test that parameters work when passed as input values."""
 
@@ -254,6 +275,7 @@ def rx_theta(theta: float):
     @aq.main
     def parametric():
         rx_theta(FreeParameter("theta"))
+        measure(0)
 
     measurements = _test_parametric_on_local_sim(parametric(), {"theta": 3.14})
     assert 0 not in measurements["__bit_0__"]

From b7c11e2311c842010e7ace361ddca6942fce560f Mon Sep 17 00:00:00 2001
From: Lauren Capelluto <laurecap@amazon.com>
Date: Fri, 27 Oct 2023 13:25:11 -0400
Subject: [PATCH 06/10] Scope out conditions

---
 .../experimental/autoqasm/test_parameters.py  | 23 -------------------
 1 file changed, 23 deletions(-)

diff --git a/test/unit_tests/braket/experimental/autoqasm/test_parameters.py b/test/unit_tests/braket/experimental/autoqasm/test_parameters.py
index 0d4f9dc90..58f0cf700 100644
--- a/test/unit_tests/braket/experimental/autoqasm/test_parameters.py
+++ b/test/unit_tests/braket/experimental/autoqasm/test_parameters.py
@@ -281,29 +281,6 @@ def parametric():
     assert 0 not in measurements["__bit_0__"]
 
 
-def test_parameter_as_condition():
-    """Test parameters used in conditional statements."""
-
-    @aq.main
-    def parametric(val: float):
-        threshold = 0.9
-        if val > threshold:
-            x(0)
-        measure(0)
-
-    expected = """OPENQASM 3.0;
-input float[64] val;
-float[64] threshold;
-qubit[1] __qubits__;
-threshold = 0.9;
-if (val > threshold) {
-    x __qubits__[0];
-}
-bit __bit_0__;
-__bit_0__ = measure __qubits__[0];"""
-    assert parametric(FreeParameter("val")).to_ir() == expected
-
-
 def test_parametric_gate_args():
     """Test that gates can be used with parameters."""
 

From d65f91bf3b7e4ab291e678aefbd73280379144f7 Mon Sep 17 00:00:00 2001
From: Lauren Capelluto <laurecap@amazon.com>
Date: Fri, 27 Oct 2023 15:08:13 -0400
Subject: [PATCH 07/10] Fix lint

---
 .../experimental/autoqasm/test_parameters.py       | 14 +-------------
 1 file changed, 1 insertion(+), 13 deletions(-)

diff --git a/test/unit_tests/braket/experimental/autoqasm/test_parameters.py b/test/unit_tests/braket/experimental/autoqasm/test_parameters.py
index 58f0cf700..fa80c6d79 100644
--- a/test/unit_tests/braket/experimental/autoqasm/test_parameters.py
+++ b/test/unit_tests/braket/experimental/autoqasm/test_parameters.py
@@ -13,23 +13,11 @@
 
 """AutoQASM tests for parameter support."""
 
-import pytest
-
 import braket.experimental.autoqasm as aq
 from braket.circuits import FreeParameter
 from braket.default_simulator import StateVectorSimulator
 from braket.devices.local_simulator import LocalSimulator
-from braket.experimental.autoqasm.instructions import (
-    cnot,
-    cphaseshift,
-    gpi,
-    h,
-    measure,
-    ms,
-    rx,
-    rz,
-    x,
-)
+from braket.experimental.autoqasm.instructions import cnot, cphaseshift, measure, ms, rx, rz
 from braket.tasks.local_quantum_task import LocalQuantumTask
 
 

From e90fd535d738b57cf2071282c370eca75e9e7d95 Mon Sep 17 00:00:00 2001
From: Lauren Capelluto <laurecap@amazon.com>
Date: Fri, 27 Oct 2023 16:26:34 -0400
Subject: [PATCH 08/10] Replace type | type with Union

---
 .../autoqasm/instructions/gates.py            | 77 ++++++++++---------
 1 file changed, 39 insertions(+), 38 deletions(-)

diff --git a/src/braket/experimental/autoqasm/instructions/gates.py b/src/braket/experimental/autoqasm/instructions/gates.py
index 13b5ae13d..40ee7518f 100644
--- a/src/braket/experimental/autoqasm/instructions/gates.py
+++ b/src/braket/experimental/autoqasm/instructions/gates.py
@@ -12,8 +12,9 @@
 # language governing permissions and limitations under the License.
 
 
-"""Quantum gates, unitary instructions, that apply to qubits.
-"""
+"""Quantum gates, unitary instructions, that apply to qubits."""
+
+from typing import Union
 
 from braket.circuits.free_parameter_expression import FreeParameterExpression
 
@@ -54,14 +55,14 @@ def cnot(
 def cphaseshift(
     control: QubitIdentifierType,
     target: QubitIdentifierType,
-    angle: float | FreeParameterExpression,
+    angle: Union[float, FreeParameterExpression],
 ) -> None:
     """Controlled phase shift gate.
 
     Args:
         control (QubitIdentifierType): Control qubit.
         target (QubitIdentifierType): Target qubit.
-        angle (float | FreeParameterExpression): Rotation angle in radians.
+        angle (Union[float, FreeParameterExpression]): Rotation angle in radians.
 
     """
     _qubit_instruction("cphaseshift", [control, target], angle)
@@ -70,14 +71,14 @@ def cphaseshift(
 def cphaseshift00(
     control: QubitIdentifierType,
     target: QubitIdentifierType,
-    angle: float | FreeParameterExpression,
+    angle: Union[float, FreeParameterExpression],
 ) -> None:
     """Controlled phase shift gate for phasing the \\|00> state.
 
     Args:
         control (QubitIdentifierType): Control qubit.
         target (QubitIdentifierType): Target qubit.
-        angle (float | FreeParameterExpression): Rotation angle in radians.
+        angle (Union[float, FreeParameterExpression]): Rotation angle in radians.
 
     """
     _qubit_instruction("cphaseshift00", [control, target], angle)
@@ -86,14 +87,14 @@ def cphaseshift00(
 def cphaseshift01(
     control: QubitIdentifierType,
     target: QubitIdentifierType,
-    angle: float | FreeParameterExpression,
+    angle: Union[float, FreeParameterExpression],
 ) -> None:
     """Controlled phase shift gate for phasing the \\|01> state.
 
     Args:
         control (QubitIdentifierType): Control qubit.
         target (QubitIdentifierType): Target qubit.
-        angle (float | FreeParameterExpression): Rotation angle in radians.
+        angle (Union[float, FreeParameterExpression]): Rotation angle in radians.
 
     """
     _qubit_instruction("cphaseshift01", [control, target], angle)
@@ -102,14 +103,14 @@ def cphaseshift01(
 def cphaseshift10(
     control: QubitIdentifierType,
     target: QubitIdentifierType,
-    angle: float | FreeParameterExpression,
+    angle: Union[float, FreeParameterExpression],
 ) -> None:
     """Controlled phase shift gate for phasing the \\|10> state.
 
     Args:
         control (QubitIdentifierType): Control qubit.
         target (QubitIdentifierType): Target qubit.
-        angle (float | FreeParameterExpression): Rotation angle in radians.
+        angle (Union[float, FreeParameterExpression]): Rotation angle in radians.
 
     """
     _qubit_instruction("cphaseshift10", [control, target], angle)
@@ -189,13 +190,13 @@ def ecr(
 
 def gpi(
     target: QubitIdentifierType,
-    angle: float | FreeParameterExpression,
+    angle: Union[float, FreeParameterExpression],
 ) -> None:
     """IonQ GPi gate.
 
     Args:
         target (QubitIdentifierType): Target qubit.
-        angle (float | FreeParameterExpression): Rotation angle in radians.
+        angle (Union[float, FreeParameterExpression]): Rotation angle in radians.
 
     """
     _qubit_instruction("gpi", [target], angle)
@@ -203,13 +204,13 @@ def gpi(
 
 def gpi2(
     target: QubitIdentifierType,
-    angle: float | FreeParameterExpression,
+    angle: Union[float, FreeParameterExpression],
 ) -> None:
     """IonQ GPi2 gate.
 
     Args:
         target (QubitIdentifierType): Target qubit.
-        angle (float | FreeParameterExpression): Rotation angle in radians.
+        angle (Union[float, FreeParameterExpression]): Rotation angle in radians.
 
     """
     _qubit_instruction("gpi2", [target], angle)
@@ -256,18 +257,18 @@ def iswap(
 def ms(
     target_0: QubitIdentifierType,
     target_1: QubitIdentifierType,
-    angle_0: float | FreeParameterExpression,
-    angle_1: float | FreeParameterExpression,
-    angle_2: float | FreeParameterExpression,
+    angle_0: Union[float, FreeParameterExpression],
+    angle_1: Union[float, FreeParameterExpression],
+    angle_2: Union[float, FreeParameterExpression],
 ) -> None:
     """IonQ Mølmer-Sørenson gate.
 
     Args:
         target_0 (QubitIdentifierType): Target qubit 0.
         target_1 (QubitIdentifierType): Target qubit 1.
-        angle_0 (float | FreeParameterExpression): Rotation angle 0 in radians.
-        angle_1 (float | FreeParameterExpression): Rotation angle 1 in radians.
-        angle_2 (float | FreeParameterExpression): Rotation angle 2 in radians.
+        angle_0 (Union[float, FreeParameterExpression]): Rotation angle 0 in radians.
+        angle_1 (Union[float, FreeParameterExpression]): Rotation angle 1 in radians.
+        angle_2 (Union[float, FreeParameterExpression]): Rotation angle 2 in radians.
 
     """
     _qubit_instruction("ms", [target_0, target_1], angle_0, angle_1, angle_2)
@@ -275,13 +276,13 @@ def ms(
 
 def phaseshift(
     target: QubitIdentifierType,
-    angle: float | FreeParameterExpression,
+    angle: Union[float, FreeParameterExpression],
 ) -> None:
     """Phase shift gate.
 
     Args:
         target (QubitIdentifierType): Target qubit.
-        angle (float | FreeParameterExpression): Rotation angle in radians.
+        angle (Union[float, FreeParameterExpression]): Rotation angle in radians.
 
     """
     _qubit_instruction("phaseshift", [target], angle)
@@ -290,14 +291,14 @@ def phaseshift(
 def pswap(
     target_0: QubitIdentifierType,
     target_1: QubitIdentifierType,
-    angle: float | FreeParameterExpression,
+    angle: Union[float, FreeParameterExpression],
 ) -> None:
     """PSwap gate.
 
     Args:
         target_0 (QubitIdentifierType): Target qubit 0.
         target_1 (QubitIdentifierType): Target qubit 1.
-        angle (float | FreeParameterExpression): Rotation angle in radians.
+        angle (Union[float, FreeParameterExpression]): Rotation angle in radians.
 
     """
     _qubit_instruction("pswap", [target_0, target_1], angle)
@@ -305,13 +306,13 @@ def pswap(
 
 def rx(
     target: QubitIdentifierType,
-    angle: float | FreeParameterExpression,
+    angle: Union[float, FreeParameterExpression],
 ) -> None:
     """X-axis rotation gate.
 
     Args:
         target (QubitIdentifierType): Target qubit.
-        angle (float | FreeParameterExpression): Rotation angle in radians.
+        angle (Union[float, FreeParameterExpression]): Rotation angle in radians.
 
     """
     _qubit_instruction("rx", [target], angle)
@@ -319,13 +320,13 @@ def rx(
 
 def ry(
     target: QubitIdentifierType,
-    angle: float | FreeParameterExpression,
+    angle: Union[float, FreeParameterExpression],
 ) -> None:
     """Y-axis rotation gate.
 
     Args:
         target (QubitIdentifierType): Target qubit.
-        angle (float | FreeParameterExpression): Rotation angle in radians.
+        angle (Union[float, FreeParameterExpression]): Rotation angle in radians.
 
     """
     _qubit_instruction("ry", [target], angle)
@@ -333,13 +334,13 @@ def ry(
 
 def rz(
     target: QubitIdentifierType,
-    angle: float | FreeParameterExpression,
+    angle: Union[float, FreeParameterExpression],
 ) -> None:
     """Z-axis rotation gate.
 
     Args:
         target (QubitIdentifierType): Target qubit.
-        angle (float | FreeParameterExpression): Rotation angle in radians.
+        angle (Union[float, FreeParameterExpression]): Rotation angle in radians.
 
     """
     _qubit_instruction("rz", [target], angle)
@@ -446,14 +447,14 @@ def x(
 def xx(
     target_0: QubitIdentifierType,
     target_1: QubitIdentifierType,
-    angle: float | FreeParameterExpression,
+    angle: Union[float, FreeParameterExpression],
 ) -> None:
     """Ising XX coupling gate.
 
     Args:
         target_0 (QubitIdentifierType): Target qubit 0.
         target_1 (QubitIdentifierType): Target qubit 1.
-        angle (float | FreeParameterExpression): Rotation angle in radians.
+        angle (Union[float, FreeParameterExpression]): Rotation angle in radians.
 
     """
     _qubit_instruction("xx", [target_0, target_1], angle)
@@ -462,14 +463,14 @@ def xx(
 def xy(
     target_0: QubitIdentifierType,
     target_1: QubitIdentifierType,
-    angle: float | FreeParameterExpression,
+    angle: Union[float, FreeParameterExpression],
 ) -> None:
     """XY gates
 
     Args:
         target_0 (QubitIdentifierType): Target qubit 0.
         target_1 (QubitIdentifierType): Target qubit 1.
-        angle (float | FreeParameterExpression): Rotation angle in radians.
+        angle (Union[float, FreeParameterExpression]): Rotation angle in radians.
 
     """
     _qubit_instruction("xy", [target_0, target_1], angle)
@@ -490,14 +491,14 @@ def y(
 def yy(
     target_0: QubitIdentifierType,
     target_1: QubitIdentifierType,
-    angle: float | FreeParameterExpression,
+    angle: Union[float, FreeParameterExpression],
 ) -> None:
     """Ising YY coupling gate.
 
     Args:
         target_0 (QubitIdentifierType): Target qubit 0.
         target_1 (QubitIdentifierType): Target qubit 1.
-        angle (float | FreeParameterExpression): Rotation angle in radians.
+        angle (Union[float, FreeParameterExpression]): Rotation angle in radians.
 
     """
     _qubit_instruction("yy", [target_0, target_1], angle)
@@ -518,14 +519,14 @@ def z(
 def zz(
     target_0: QubitIdentifierType,
     target_1: QubitIdentifierType,
-    angle: float | FreeParameterExpression,
+    angle: Union[float, FreeParameterExpression],
 ) -> None:
     """Ising ZZ coupling gate.
 
     Args:
         target_0 (QubitIdentifierType): Target qubit 0.
         target_1 (QubitIdentifierType): Target qubit 1.
-        angle (float | FreeParameterExpression): Rotation angle in radians.
+        angle (Union[float, FreeParameterExpression]): Rotation angle in radians.
 
     """
     _qubit_instruction("zz", [target_0, target_1], angle)

From 0f1f1a2353c844fcbc6a99b44ee361817fe94329 Mon Sep 17 00:00:00 2001
From: Lauren Capelluto <laurecap@amazon.com>
Date: Sat, 28 Oct 2023 12:36:35 -0400
Subject: [PATCH 09/10] Fix coverage, remove dup test

---
 src/braket/experimental/autoqasm/api.py       |  5 ++--
 .../experimental/autoqasm/test_parameters.py  | 23 -------------------
 2 files changed, 2 insertions(+), 26 deletions(-)

diff --git a/src/braket/experimental/autoqasm/api.py b/src/braket/experimental/autoqasm/api.py
index 2acea01bb..6e3c4feef 100644
--- a/src/braket/experimental/autoqasm/api.py
+++ b/src/braket/experimental/autoqasm/api.py
@@ -547,9 +547,8 @@ def _get_gate_args(f: Callable) -> aq_program.GateArgs:
 
         if param.annotation == aq_instructions.QubitIdentifierType:
             gate_args.append_qubit(param.name)
-        elif param.annotation in [float, aq_types.FloatVar] or (
-            get_args(param.annotation)
-            and any(type_ in [float, aq_types.FloatVar] for type_ in get_args(param.annotation))
+        elif param.annotation == float or any(
+            type_ == float for type_ in get_args(param.annotation)
         ):
             gate_args.append_angle(param.name)
         else:
diff --git a/test/unit_tests/braket/experimental/autoqasm/test_parameters.py b/test/unit_tests/braket/experimental/autoqasm/test_parameters.py
index fa80c6d79..137327a31 100644
--- a/test/unit_tests/braket/experimental/autoqasm/test_parameters.py
+++ b/test/unit_tests/braket/experimental/autoqasm/test_parameters.py
@@ -78,29 +78,6 @@ def test_multiple_parameters():
     assert multi_parametric().to_ir() == expected
 
 
-def test_typed_parameters():
-    """Test that multiple free parameters all appear in the processed program."""
-
-    @aq.main
-    def multi_parametric():
-        rx(0, FreeParameter("alpha"))
-        rx(1, FreeParameter("theta"))
-        c = measure([0, 1])  # noqa: F841
-
-    expected = """OPENQASM 3.0;
-bit[2] c;
-input float[64] alpha;
-input float[64] theta;
-qubit[2] __qubits__;
-rx(alpha) __qubits__[0];
-rx(theta) __qubits__[1];
-bit[2] __bit_0__ = "00";
-__bit_0__[0] = measure __qubits__[0];
-__bit_0__[1] = measure __qubits__[1];
-c = __bit_0__;"""
-    assert multi_parametric().to_ir() == expected
-
-
 def test_sim_multi_param():
     measurements = _test_parametric_on_local_sim(multi_parametric(), {"alpha": 3.14, "theta": 0})
     assert all(val == "10" for val in measurements["c"])

From c4ae77aecaa3045a0ed344afa2c457d6cd72c8fa Mon Sep 17 00:00:00 2001
From: Lauren Capelluto <laurecap@amazon.com>
Date: Mon, 30 Oct 2023 11:33:34 -0400
Subject: [PATCH 10/10] Add test for gate calibrations

---
 .../experimental/autoqasm/test_parameters.py  | 22 +++++++++++++++++++
 1 file changed, 22 insertions(+)

diff --git a/test/unit_tests/braket/experimental/autoqasm/test_parameters.py b/test/unit_tests/braket/experimental/autoqasm/test_parameters.py
index 137327a31..3971d45c2 100644
--- a/test/unit_tests/braket/experimental/autoqasm/test_parameters.py
+++ b/test/unit_tests/braket/experimental/autoqasm/test_parameters.py
@@ -17,6 +17,7 @@
 from braket.circuits import FreeParameter
 from braket.default_simulator import StateVectorSimulator
 from braket.devices.local_simulator import LocalSimulator
+from braket.experimental.autoqasm import pulse
 from braket.experimental.autoqasm.instructions import cnot, cphaseshift, measure, ms, rx, rz
 from braket.tasks.local_quantum_task import LocalQuantumTask
 
@@ -265,3 +266,24 @@ def parametric():
 qubit[3] __qubits__;
 rx_theta(θ) __qubits__[2];"""
     assert parametric().to_ir() == expected
+
+
+def test_parametric_pulse_cals():
+    """Test that pulse calibrations work with free parameters."""
+
+    @aq.gate_calibration(implements=rx, target="$1")
+    def cal_1(angle: float):
+        pulse.delay("$1", angle)
+
+    @aq.main
+    def my_program():
+        rx("$1", FreeParameter("theta"))
+
+    expected = """OPENQASM 3.0;
+defcal rx(angle[32] angle) $1 {
+    delay[(angle) * 1s] $1;
+}
+input float[64] theta;
+rx(theta) $1;"""
+    qasm = my_program().with_calibrations(cal_1).to_ir()
+    assert qasm == expected