Respond to CR

src/braket/experimental/autoqasm/errors.py

@@ -37,6 +37,10 @@ class MissingParameterTypeError(AutoQasmError):
     """AutoQASM requires type hints for subroutine parameters."""
 
 
+class ParameterNotFoundError(AutoQasmError):
+    """A FreeParameter could not be found in the program."""
+
+
 class InvalidGateDefinition(AutoQasmError):
     """Gate definition does not meet the necessary requirements."""
 

src/braket/experimental/autoqasm/operators/utils.py

@@ -14,27 +14,29 @@
 
 "Utility methods for operators."
 
-from typing import Any
+from typing import Any, Union
 
 from braket.circuits import FreeParameter
 from braket.experimental.autoqasm import program
 from braket.experimental.autoqasm import types as aq_types
 
 
-def _register_and_convert_parameters(*args: list[Any]) -> list[aq_types.FloatVar]:
+def _register_and_convert_parameters(
+    *args: tuple[Any],
+) -> Union[list[aq_types.FloatVar], aq_types.FloatVar]:
     """Adds FreeParameters to the program conversion context parameter registry, and
     returns the associated FloatVar objects.
 
     Notes: Adding a parameter to the registry twice is safe. Conversion is a pass through
-    for non-FreeParameter inputs.
+    for non-FreeParameter inputs. Input and output arity is the same.
 
     FloatVars are more compatible with the program conversion operations.
 
     Returns:
-        list[FloatVar]: FloatVars for program conversion.
+        Union[list[FloatVar], FloatVar]: FloatVars for program conversion.
     """
     program_conversion_context = program.get_program_conversion_context()
-    program_conversion_context.register_args(args)  # TODO could be one item
+    program_conversion_context.register_args(args)
     result = []
     for arg in args:
         if isinstance(arg, FreeParameter):

src/braket/experimental/autoqasm/program/program.py

@@ -130,11 +130,12 @@ def make_bound_program(self, param_values: dict[str, float], strict: bool = Fals
         Args:
             param_values (dict[str, float]): A mapping of FreeParameter names
                 to a value to assign to them.
-            strict (bool): If True, raises a ValueError if any of the FreeParameters
+            strict (bool): If True, raises a ParameterNotFoundError if any of the FreeParameters
                 in param_values do not appear in the program. False by default.
 
         Raises:
-            ValueError: If a parameter name is given which does not appear in the program.
+            ParameterNotFoundError: If a parameter name is given which does not appear in
+                the program.
 
         Returns:
             Program: Returns a program with all present parameters fixed to their respective
@@ -148,7 +149,7 @@ def make_bound_program(self, param_values: dict[str, float], strict: bool = Fals
                 assert target.init_expression == "input", "Only free parameters can be bound."
                 target.init_expression = value
             elif strict:
-                raise ValueError(f"No parameter in the program named: {name}")
+                raise errors.ParameterNotFoundError(f"No parameter in the program named: {name}")
 
         return Program(bound_oqpy_program, self._has_pulse_control)
 
@@ -345,13 +346,14 @@ def get_parameter(self, name: str) -> oqpy.FloatVar:
             name (str): The name of the parameter.
 
         Raises:
-            ValueError: If there is no parameter with the given name registered with the program.
+            ParameterNotFoundError: If there is no parameter with the given name registered
+            with the program.
 
         Returns:
             FloatVar: The associated variable.
         """
         if name not in self._free_parameters:
-            raise ValueError(f"Free parameter '{name}' was not found.")
+            raise errors.ParameterNotFoundError(f"Free parameter '{name}' was not found.")
         return self._free_parameters[name]
 
     def get_free_parameters(self) -> list[oqpy.FloatVar]:

src/braket/experimental/autoqasm/transpiler/transpiler.py

@@ -144,14 +144,14 @@ def transform_ast(
         # canonicalization creates.
         node = continue_statements.transform(node, ctx)
         node = return_statements.transform(node, ctx)
-        node = comparisons.transform(node, ctx)
         node = assignments.transform(node, ctx)
         node = lists.transform(node, ctx)
         node = slices.transform(node, ctx)
         node = call_trees.transform(node, ctx)
         node = control_flow.transform(node, ctx)
         node = conditional_expressions.transform(node, ctx)
         node = logical_expressions.transform(node, ctx)
+        node = comparisons.transform(node, ctx)
         node = variables.transform(node, ctx)
 
         return node

test/unit_tests/braket/experimental/autoqasm/test_parameters.py

@@ -461,7 +461,9 @@ def parametric(theta: float):
         measure(0)
 
     prog = parametric(FreeParameter("alpha"))
-    with pytest.raises(ValueError, match="No parameter in the program named: beta"):
+    with pytest.raises(
+        aq.errors.ParameterNotFoundError, match="No parameter in the program named: beta"
+    ):
         prog.make_bound_program({"beta": 0.5}, strict=True)
 
 
@@ -484,7 +486,9 @@ def test_binding_variable_fails():
     def parametric():
         alpha = aq.FloatVar(1.2)  # noqa: F841
 
-    with pytest.raises(ValueError, match="No parameter in the program named: beta"):
+    with pytest.raises(
+        aq.errors.ParameterNotFoundError, match="No parameter in the program named: beta"
+    ):
         parametric().make_bound_program({"beta": 0.5}, strict=True)
 
 

test/unit_tests/braket/experimental/autoqasm/test_program.py

@@ -44,7 +44,7 @@ def test_get_parameter_invalid_name():
     """Tests the get_parameter function."""
     prog = aq.program.ProgramConversionContext()
     prog.register_parameter(FreeParameter("alpha"))
-    with pytest.raises(ValueError):
+    with pytest.raises(aq.errors.ParameterNotFoundError):
         prog.get_parameter("not_a_parameter")