feature: Add support for FreeParameters in AutoQASM conditional state…

…ments (#789)

* feature: FreeParameters in conditional statements

* Update type check to include FreeParameter as a qasm type
* Add support for free parameters in logical operations
* Add support for comparison statements
---------

Co-authored-by: Ryan Shaffer <3620100+rmshaffer@users.noreply.github.com>

src/braket/experimental/autoqasm/converters/comparisons.py

@@ -0,0 +1,70 @@
+# 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.
+
+"""Converters for comparison nodes."""
+
+import ast
+
+import gast
+
+from braket.experimental.autoqasm.autograph.core import ag_ctx, converter
+from braket.experimental.autoqasm.autograph.pyct import templates
+
+COMPARISON_OPERATORS = {
+    gast.Lt: "ag__.lt_",
+    gast.LtE: "ag__.lteq_",
+    gast.Gt: "ag__.gt_",
+    gast.GtE: "ag__.gteq_",
+}
+
+
+class ComparisonTransformer(converter.Base):
+    """Transformer for comparison nodes."""
+
+    def visit_Compare(self, node: ast.stmt) -> ast.stmt:
+        """Transforms a comparison node.
+
+        Args:
+            node (ast.stmt): AST node to transform.
+
+        Returns:
+            ast.stmt: Transformed node.
+        """
+        node = self.generic_visit(node)
+
+        op_type = type(node.ops[0])
+        if op_type not in COMPARISON_OPERATORS:
+            return node
+
+        template = f"{COMPARISON_OPERATORS[op_type]}(lhs_, rhs_)"
+
+        return templates.replace(
+            template,
+            lhs_=node.left,
+            rhs_=node.comparators[0],
+            original=node,
+        )[0].value
+
+
+def transform(node: ast.stmt, ctx: ag_ctx.ControlStatusCtx) -> ast.stmt:
+    """Transform comparison nodes.
+
+    Args:
+        node (ast.stmt): AST node to transform.
+        ctx (ag_ctx.ControlStatusCtx): Transformer context.
+
+    Returns:
+        ast.stmt: Transformed node.
+    """
+
+    return ComparisonTransformer(ctx).visit(node)

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/__init__.py

@@ -27,6 +27,7 @@
 )
 
 from .assignments import assign_stmt  # noqa: F401
+from .comparisons import gt_, gteq_, lt_, lteq_  # noqa: F401
 from .conditional_expressions import if_exp  # noqa: F401
 from .control_flow import for_stmt, if_stmt, while_stmt  # noqa: F401
 from .data_structures import ListPopOpts  # noqa: F401

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

@@ -0,0 +1,126 @@
+# 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.
+
+
+"""Operators for comparison operators: <, <=, >, and >=."""
+
+from typing import Any, Union
+
+from braket.experimental.autoqasm import program
+from braket.experimental.autoqasm import types as aq_types
+
+from .utils import _register_and_convert_parameters
+
+
+def lt_(a: Any, b: Any) -> Union[bool, aq_types.BoolVar]:
+    """Functional form of "<".
+
+    Args:
+        a (Any): The first expression.
+        b (Any): The second expression.
+
+    Returns:
+        Union[bool, BoolVar]: Whether the first expression is less than the second.
+    """
+    if aq_types.is_qasm_type(a) or aq_types.is_qasm_type(b):
+        return _aq_lt(a, b)
+    else:
+        return a < b
+
+
+def _aq_lt(a: Any, b: Any) -> aq_types.BoolVar:
+    a, b = _register_and_convert_parameters(a, b)
+
+    oqpy_program = program.get_program_conversion_context().get_oqpy_program()
+    result = aq_types.BoolVar()
+    oqpy_program.declare(result)
+    oqpy_program.set(result, a < b)
+    return result
+
+
+def lteq_(a: Any, b: Any) -> Union[bool, aq_types.BoolVar]:
+    """Functional form of "<=".
+
+    Args:
+        a (Any): The first expression.
+        b (Any): The second expression.
+
+    Returns:
+        Union[bool, BoolVar]: Whether the first expression is less than or equal to the second.
+    """
+    if aq_types.is_qasm_type(a) or aq_types.is_qasm_type(b):
+        return _aq_lteq(a, b)
+    else:
+        return a <= b
+
+
+def _aq_lteq(a: Any, b: Any) -> aq_types.BoolVar:
+    a, b = _register_and_convert_parameters(a, b)
+
+    oqpy_program = program.get_program_conversion_context().get_oqpy_program()
+    result = aq_types.BoolVar()
+    oqpy_program.declare(result)
+    oqpy_program.set(result, a <= b)
+    return result
+
+
+def gt_(a: Any, b: Any) -> Union[bool, aq_types.BoolVar]:
+    """Functional form of ">".
+
+    Args:
+        a (Any): The first expression.
+        b (Any): The second expression.
+
+    Returns:
+        Union[bool, BoolVar]: Whether the first expression is greater than the second.
+    """
+    if aq_types.is_qasm_type(a) or aq_types.is_qasm_type(b):
+        return _aq_gt(a, b)
+    else:
+        return a > b
+
+
+def _aq_gt(a: Any, b: Any) -> aq_types.BoolVar:
+    a, b = _register_and_convert_parameters(a, b)
+
+    oqpy_program = program.get_program_conversion_context().get_oqpy_program()
+    result = aq_types.BoolVar()
+    oqpy_program.declare(result)
+    oqpy_program.set(result, a > b)
+    return result
+
+
+def gteq_(a: Any, b: Any) -> Union[bool, aq_types.BoolVar]:
+    """Functional form of ">=".
+
+    Args:
+        a (Any): The first expression.
+        b (Any): The second expression.
+
+    Returns:
+        Union[bool, BoolVar]: Whether the first expression is greater than or equal to the second.
+    """
+    if aq_types.is_qasm_type(a) or aq_types.is_qasm_type(b):
+        return _aq_gteq(a, b)
+    else:
+        return a >= b
+
+
+def _aq_gteq(a: Any, b: Any) -> aq_types.BoolVar:
+    a, b = _register_and_convert_parameters(a, b)
+
+    oqpy_program = program.get_program_conversion_context().get_oqpy_program()
+    result = aq_types.BoolVar()
+    oqpy_program.declare(result)
+    oqpy_program.set(result, a >= b)
+    return result

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

@@ -22,6 +22,8 @@
 from braket.experimental.autoqasm import program
 from braket.experimental.autoqasm import types as aq_types
 
+from .utils import _register_and_convert_parameters
+
 
 def and_(a: Callable[[], Any], b: Callable[[], Any]) -> Union[bool, aq_types.BoolVar]:
     """Functional form of "and".
@@ -42,6 +44,8 @@ def and_(a: Callable[[], Any], b: Callable[[], Any]) -> Union[bool, aq_types.Boo
 
 
 def _oqpy_and(a: Any, b: Any) -> aq_types.BoolVar:
+    a, b = _register_and_convert_parameters(a, b)
+
     oqpy_program = program.get_program_conversion_context().get_oqpy_program()
     result = aq_types.BoolVar()
     oqpy_program.declare(result)
@@ -72,6 +76,8 @@ def or_(a: Callable[[], Any], b: Callable[[], Any]) -> Union[bool, aq_types.Bool
 
 
 def _oqpy_or(a: Any, b: Any) -> aq_types.BoolVar:
+    a, b = _register_and_convert_parameters(a, b)
+
     oqpy_program = program.get_program_conversion_context().get_oqpy_program()
     result = aq_types.BoolVar()
     oqpy_program.declare(result)
@@ -99,6 +105,8 @@ def not_(a: Any) -> Union[bool, aq_types.BoolVar]:
 
 
 def _oqpy_not(a: Any) -> aq_types.BoolVar:
+    a = _register_and_convert_parameters(a)
+
     oqpy_program = program.get_program_conversion_context().get_oqpy_program()
     result = aq_types.BoolVar()
     oqpy_program.declare(result)
@@ -127,6 +135,8 @@ def eq(a: Any, b: Any) -> Union[bool, aq_types.BoolVar]:
 
 
 def _oqpy_eq(a: Any, b: Any) -> aq_types.BoolVar:
+    a, b = _register_and_convert_parameters(a, b)
+
     oqpy_program = program.get_program_conversion_context().get_oqpy_program()
     is_equal = aq_types.BoolVar()
     oqpy_program.declare(is_equal)
@@ -155,6 +165,8 @@ def not_eq(a: Any, b: Any) -> Union[bool, aq_types.BoolVar]:
 
 
 def _oqpy_not_eq(a: Any, b: Any) -> aq_types.BoolVar:
+    a, b = _register_and_convert_parameters(a, b)
+
     oqpy_program = program.get_program_conversion_context().get_oqpy_program()
     is_not_equal = aq_types.BoolVar()
     oqpy_program.declare(is_not_equal)

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

@@ -0,0 +1,47 @@
+# 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.
+
+
+"Utility methods for operators."
+
+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: 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. Input and output arity is the same.
+
+    FloatVars are more compatible with the program conversion operations.
+
+    Returns:
+        Union[list[FloatVar], FloatVar]: FloatVars for program conversion.
+    """
+    program_conversion_context = program.get_program_conversion_context()
+    program_conversion_context.register_args(args)
+    result = []
+    for arg in args:
+        if isinstance(arg, FreeParameter):
+            var = program.get_program_conversion_context().get_parameter(arg.name)
+            result.append(var)
+        else:
+            result.append(arg)
+    return result[0] if len(result) == 1 else result

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)
 
@@ -321,22 +322,39 @@ def register_args(self, args: list[Any]) -> None:
         """
         for arg in args:
             if isinstance(arg, FreeParameter):
-                self.register_parameter(arg.name)
+                self.register_parameter(arg)
             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.
+    def register_parameter(self, parameter: FreeParameter) -> None:
+        """Register an input parameter if it has not already been registered.
+        Only floats are currently supported.
 
         Args:
-            name (str): The identifier for the parameter.
+            parameter (FreeParameter): The parameter to register with the program.
+        """
+        if parameter.name not in self._free_parameters:
+            self._free_parameters[parameter.name] = oqpy.FloatVar("input", name=parameter.name)
+
+    def get_parameter(self, name: str) -> oqpy.FloatVar:
+        """Return a named oqpy.FloatVar that is used as a free parameter in the program.
+
+        Args:
+            name (str): The name of the parameter.
+
+        Raises:
+            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:
-            self._free_parameters[name] = oqpy.FloatVar("input", name=name)
+            raise errors.ParameterNotFoundError(f"Free parameter '{name}' was not found.")
+        return self._free_parameters[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."""