classiq 0.66.1__py3-none-any.whl → 0.68.0__py3-none-any.whl

classiq/model_expansions/quantum_operations/block_evaluator.py

@@ -0,0 +1,76 @@
+from collections.abc import Sequence
+from typing import TYPE_CHECKING
+
+from classiq.interface.model.quantum_statement import QuantumOperation, QuantumStatement
+
+from classiq.model_expansions.closure import Closure
+from classiq.model_expansions.quantum_operations.emitter import Emitter
+from classiq.model_expansions.scope import Scope
+
+if TYPE_CHECKING:
+    from classiq.model_expansions.interpreters.base_interpreter import BaseInterpreter
+
+
+_BLOCK_RENAMES = {
+    "compute": "within",
+    "action": "apply",
+}
+_REVERSE_BLOCK_RENAMES = {rename: name for name, rename in _BLOCK_RENAMES.items()}
+
+
+class BlockEvaluator(Emitter[QuantumOperation]):
+    def __init__(
+        self, interpreter: "BaseInterpreter", operation_name: str, *block_names: str
+    ) -> None:
+        super().__init__(interpreter)
+        self._operation_name = operation_name
+        self._block_names: Sequence[str] = block_names
+
+    def emit(self, op: QuantumOperation, /) -> bool:
+        expanded_blocks: dict[str, list[QuantumStatement]] = {}
+        blocks = [
+            block
+            for block in self._block_names
+            if hasattr(op, block) and getattr(op, block) is not None
+        ]
+
+        if len(blocks) > 0:
+            if op.is_generative():
+                expanded_blocks = self.expand_generative_blocks(op)
+            else:
+                expanded_blocks = self.expand_blocks(op, blocks)
+            expanded_blocks.update(expanded_blocks)
+
+        op = op.model_copy(update={**expanded_blocks, "back_ref": op.uuid})
+        self._builder.emit_statement(op)
+        return True
+
+    def expand_blocks(
+        self, op: QuantumOperation, block_names: list[str]
+    ) -> dict[str, list[QuantumStatement]]:
+        blocks = {
+            _BLOCK_RENAMES.get(block, block): getattr(op, block)
+            for block in block_names
+        }
+        block_closure = Closure(
+            name=self._operation_name,
+            scope=Scope(parent=self._current_scope),
+            blocks=blocks,
+        )
+        context = self._expand_operation(block_closure)
+        return {
+            block: context.statements(_BLOCK_RENAMES.get(block, block))
+            for block in block_names
+        }
+
+    def expand_generative_blocks(
+        self, op: QuantumOperation
+    ) -> dict[str, list[QuantumStatement]]:
+        blocks = [
+            block for block in self._block_names if op.has_generative_block(block)
+        ]
+        context = self._expand_generative_context(op, self._operation_name, blocks)
+        return {
+            _REVERSE_BLOCK_RENAMES.get(block, block): context.statements(block)
+            for block in blocks
+        }

classiq/model_expansions/quantum_operations/call_emitter.py

@@ -1,13 +1,17 @@
 from collections.abc import Sequence
-from itertools import chain
+from itertools import chain, combinations
 from typing import (
     TYPE_CHECKING,
     Generic,
-    Optional,
     cast,
 )
+from uuid import UUID
 
 from classiq.interface.debug_info.debug_info import FunctionDebugInfo
+from classiq.interface.exceptions import ClassiqExpansionError
+from classiq.interface.generator.functions.port_declaration import (
+    PortDeclarationDirection,
+)
 from classiq.interface.generator.generated_circuit_data import OperationLevel
 from classiq.interface.model.classical_parameter_declaration import (
     ClassicalParameterDeclaration,
@@ -26,6 +30,8 @@ from classiq.interface.model.variable_declaration_statement import (
 )
 
 from classiq.model_expansions.capturing.captured_vars import (
+    INITIALIZED_VAR_MESSAGE,
+    UNINITIALIZED_VAR_MESSAGE,
     validate_args_are_not_propagated,
 )
 from classiq.model_expansions.closure import FunctionClosure
@@ -53,6 +59,25 @@ if TYPE_CHECKING:
     from classiq.model_expansions.interpreters.base_interpreter import BaseInterpreter
 
 
+def _validate_cloning(evaluated_args: list[Evaluated]) -> None:
+    handles = [
+        arg.value.handle
+        for arg in evaluated_args
+        if isinstance(arg.value, QuantumSymbol)
+    ]
+    for handle, other_handle in combinations(handles, 2):
+        if handle.overlaps(other_handle):
+            if handle == other_handle:
+                raise ClassiqExpansionError(
+                    f"Quantum cloning violation: Argument {str(handle)!r} is "
+                    f"duplicated"
+                )
+            raise ClassiqExpansionError(
+                f"Quantum cloning violation: Arguments {str(handle)!r} and "
+                f"{str(other_handle)!r} overlap"
+            )
+
+
 class CallEmitter(Generic[QuantumStatementT], Emitter[QuantumStatementT], VarSplitter):
     def __init__(self, interpreter: "BaseInterpreter") -> None:
         Emitter.__init__(self, interpreter)
@@ -71,15 +96,15 @@ class CallEmitter(Generic[QuantumStatementT], Emitter[QuantumStatementT], VarSpl
             name=self._counted_name_allocator.allocate(name),
             body=body,
             scope=Scope(parent=self._current_scope),
-            is_lambda=True,
+            lambda_external_vars=self._builder.current_block.captured_vars,
         )
-        return self._create_quantum_function_call(wrapping_function, list())
+        return self._create_quantum_function_call(wrapping_function, list(), None)
 
     def _emit_quantum_function_call(
         self,
         function: FunctionClosure,
         args: list[ArgValue],
-        propagated_debug_info: Optional[FunctionDebugInfo] = None,
+        propagated_debug_info: FunctionDebugInfo | None,
     ) -> QuantumFunctionCall:
         call = self._create_quantum_function_call(
             function, args, propagated_debug_info=propagated_debug_info
@@ -87,15 +112,27 @@ class CallEmitter(Generic[QuantumStatementT], Emitter[QuantumStatementT], VarSpl
         self.emit_statement(call)
         return call
 
+    @staticmethod
+    def _get_back_ref(
+        propagated_debug_info: FunctionDebugInfo | None,
+    ) -> UUID | None:
+        if propagated_debug_info is None:
+            return None
+        if propagated_debug_info.node is None:
+            return None
+        return propagated_debug_info.node.uuid
+
     def _create_quantum_function_call(
         self,
         function: FunctionClosure,
         args: list[ArgValue],
-        propagated_debug_info: Optional[FunctionDebugInfo] = None,
+        propagated_debug_info: FunctionDebugInfo | None,
     ) -> QuantumFunctionCall:
         function = function.clone()
         function = function.set_depth(self._builder.current_function.depth + 1)
+        self._validate_call_args(function.positional_arg_declarations, args)
         evaluated_args = [self._interpreter.evaluate(arg) for arg in args]
+        _validate_cloning(evaluated_args)
         new_declaration = self._prepare_fully_typed_declaration(
             function, evaluated_args
         )
@@ -114,13 +151,15 @@ class CallEmitter(Generic[QuantumStatementT], Emitter[QuantumStatementT], VarSpl
         new_positional_args = self._get_new_positional_args(
             evaluated_args, is_atomic, new_positional_arg_decls
         )
-        captured_args = function.captured_vars.get_captured_args(
+        captured_args = function.captured_vars.filter_vars(function).get_captured_args(
             self._builder.current_function
         )
         validate_args_are_not_propagated(new_positional_args, captured_args)
         new_positional_args.extend(captured_args)
         new_call = QuantumFunctionCall(
-            function=new_declaration.name, positional_args=new_positional_args
+            function=new_declaration.name,
+            positional_args=new_positional_args,
+            back_ref=self._get_back_ref(propagated_debug_info),
         )
         is_allocate_or_free = new_call.func_name == free.func_decl.name
         parameters = {
@@ -144,6 +183,7 @@ class CallEmitter(Generic[QuantumStatementT], Emitter[QuantumStatementT], VarSpl
             ),
             is_allocate_or_free=is_allocate_or_free,
             port_to_passed_variable_map=port_to_passed_variable_map,
+            node=new_call._as_back_ref(),
         )
         new_call.set_func_decl(new_declaration)
         return new_call
@@ -187,7 +227,10 @@ class CallEmitter(Generic[QuantumStatementT], Emitter[QuantumStatementT], VarSpl
     ) -> NativeFunctionDefinition:
         func_def = self._builder.create_definition(function_context)
 
-        captured_ports = function_context.closure.captured_vars.get_captured_ports()
+        captured_vars = function_context.closure.captured_vars.filter_vars(
+            function_context.closure
+        )
+        captured_ports = captured_vars.get_captured_ports()
         if len(captured_ports) == 0:
             return func_def
         func_def.positional_arg_declarations = list(
@@ -197,7 +240,7 @@ class CallEmitter(Generic[QuantumStatementT], Emitter[QuantumStatementT], VarSpl
         if not function_context.is_lambda:
             return func_def
         func_def.body = self.rewrite(
-            func_def.body, function_context.closure.captured_vars.get_captured_mapping()
+            func_def.body, captured_vars.get_captured_mapping()
         )
 
         return func_def
@@ -265,3 +308,29 @@ class CallEmitter(Generic[QuantumStatementT], Emitter[QuantumStatementT], VarSpl
                 evaluated_args,
             ),
         )
+
+    def _validate_call_args(
+        self, params: Sequence[PositionalArg], args: list[ArgValue]
+    ) -> None:
+        for param, arg in zip(params, args):
+            if not isinstance(param, PortDeclaration) or not isinstance(
+                arg, HandleBinding
+            ):
+                continue
+            var_name = arg.name
+            symbol = self._interpreter.evaluate(var_name)
+            if (
+                not isinstance(symbol.value, QuantumSymbol)
+                or symbol.defining_function is None
+            ):
+                continue
+            var_state = self._builder.current_block.captured_vars.get_state(
+                var_name, symbol.defining_function
+            )
+            if not var_state and param.direction in (
+                PortDeclarationDirection.Inout,
+                PortDeclarationDirection.Input,
+            ):
+                raise ClassiqExpansionError(UNINITIALIZED_VAR_MESSAGE.format(var_name))
+            if var_state and param.direction == PortDeclarationDirection.Output:
+                raise ClassiqExpansionError(INITIALIZED_VAR_MESSAGE.format(var_name))

classiq/model_expansions/quantum_operations/classicalif.py

@@ -1,5 +1,6 @@
 from collections.abc import Sequence
 
+from classiq.interface.debug_info.debug_info import new_function_debug_info_by_node
 from classiq.interface.model.classical_if import ClassicalIf
 from classiq.interface.model.quantum_function_call import QuantumFunctionCall
 from classiq.interface.model.quantum_statement import QuantumStatement
@@ -17,7 +18,7 @@ def _is_all_identity_calls(body: Sequence[QuantumStatement]) -> bool:
 
 
 class ClassicalIfEmitter(CallEmitter[ClassicalIf]):
-    def emit(self, classical_if: ClassicalIf, /) -> None:
+    def emit(self, classical_if: ClassicalIf, /) -> bool:
         condition = self._interpreter.evaluate(classical_if.condition).as_type(bool)
         op_name = "then" if condition else "else"
         is_generative = classical_if.is_generative()
@@ -25,7 +26,7 @@ class ClassicalIfEmitter(CallEmitter[ClassicalIf]):
         body: Sequence[QuantumStatement]
         if is_generative:
             if not classical_if.has_generative_block(op_name):
-                return
+                return True
             context = self._expand_generative_context(classical_if, op_name, op_name)
             context.blocks["body"] = context.blocks[op_name]
             context.blocks.pop(op_name)
@@ -34,7 +35,7 @@ class ClassicalIfEmitter(CallEmitter[ClassicalIf]):
             body = classical_if.then if condition else classical_if.else_
 
         if _is_all_identity_calls(body):
-            return
+            return True
 
         if is_generative or not self._should_wrap(body):
             for stmt in body:
@@ -42,12 +43,15 @@ class ClassicalIfEmitter(CallEmitter[ClassicalIf]):
                     self._interpreter._builder.emit_statement(stmt)
                 else:
                     self._interpreter.emit_statement(stmt)
-            return
+            return True
 
         then_else_func = FunctionClosure.create(
             name=self._counted_name_allocator.allocate("then" if condition else "else"),
             body=body,
             scope=Scope(parent=self._current_scope),
-            is_lambda=True,
+            lambda_external_vars=self._builder.current_block.captured_vars,
         )
-        self._emit_quantum_function_call(then_else_func, list())
+        self._emit_quantum_function_call(
+            then_else_func, list(), new_function_debug_info_by_node(classical_if)
+        )
+        return True

classiq/model_expansions/quantum_operations/composite_emitter.py

@@ -0,0 +1,27 @@
+from collections.abc import Sequence
+from typing import TYPE_CHECKING, Generic
+
+from classiq.model_expansions.quantum_operations.emitter import (
+    Emitter,
+    QuantumStatementT,
+)
+
+if TYPE_CHECKING:
+    from classiq.model_expansions.interpreters.base_interpreter import BaseInterpreter
+
+
+class CompositeEmitter(Generic[QuantumStatementT], Emitter[QuantumStatementT]):
+    def __init__(
+        self,
+        interpreter: "BaseInterpreter",
+        emitters: Sequence[Emitter[QuantumStatementT]],
+    ) -> None:
+        super().__init__(interpreter)
+        self._emitters = emitters
+
+    def emit(self, statement: QuantumStatementT, /) -> bool:
+        for emitter in self._emitters:
+            if emitter.emit(statement):
+                return True
+        self._builder.emit_statement(statement)
+        return True

classiq/model_expansions/quantum_operations/emitter.py

@@ -1,3 +1,4 @@
+import ast
 from abc import ABC, abstractmethod
 from typing import (
     TYPE_CHECKING,
@@ -9,7 +10,10 @@ from typing import (
 
 import sympy
 
-from classiq.interface.debug_info.debug_info import DebugInfoCollection
+from classiq.interface.debug_info.debug_info import (
+    DebugInfoCollection,
+    new_function_debug_info_by_node,
+)
 from classiq.interface.exceptions import ClassiqInternalExpansionError
 from classiq.interface.generator.expressions.evaluated_expression import (
     EvaluatedExpression,
@@ -36,12 +40,15 @@ from classiq.model_expansions.sympy_conversion.sympy_to_python import (
     translate_sympy_quantum_expression,
 )
 from classiq.model_expansions.utils.counted_name_allocator import CountedNameAllocator
+from classiq.model_expansions.visitors.variable_references import VarRefCollector
 from classiq.qmod.quantum_function import GenerativeQFunc
 
 if TYPE_CHECKING:
     from classiq.model_expansions.interpreters.base_interpreter import BaseInterpreter
 
-QuantumStatementT = TypeVar("QuantumStatementT", bound=QuantumStatement)
+QuantumStatementT = TypeVar(
+    "QuantumStatementT", bound=QuantumStatement, contravariant=True
+)
 
 
 class Emitter(Generic[QuantumStatementT], ABC):
@@ -57,7 +64,7 @@ class Emitter(Generic[QuantumStatementT], ABC):
         )
 
     @abstractmethod
-    def emit(self, statement: QuantumStatementT, /) -> None:
+    def emit(self, statement: QuantumStatementT, /) -> bool:
         pass
 
     def _expand_operation(self, closure: Closure) -> OperationContext:
@@ -140,6 +147,10 @@ class Emitter(Generic[QuantumStatementT], ABC):
     def emit_statement(self, statement: QuantumStatement) -> None:
         if isinstance(statement, QuantumOperation):
             self._update_captured_vars(statement)
+        if statement.uuid not in self._interpreter._model.debug_info:
+            self._interpreter._model.debug_info[statement.uuid] = (
+                new_function_debug_info_by_node(statement)  # type:ignore[arg-type]
+            )
         self._builder.emit_statement(statement)
 
     def _update_captured_vars(self, op: QuantumOperation) -> None:
@@ -166,3 +177,13 @@ class Emitter(Generic[QuantumStatementT], ABC):
             defining_function=defining_function,
             direction=direction,
         )
+
+    def _get_symbols_in_expression(self, expr: Expression) -> list[QuantumSymbol]:
+        vrc = VarRefCollector(ignore_duplicated_handles=True)
+        vrc.visit(ast.parse(expr.expr))
+        handles = dict.fromkeys(
+            handle
+            for handle in vrc.var_handles
+            if isinstance(self._current_scope[handle.name].value, QuantumSymbol)
+        )
+        return [self._interpreter.evaluate(handle).value for handle in handles]

classiq/model_expansions/quantum_operations/expression_evaluator.py

@@ -0,0 +1,33 @@
+from typing import TYPE_CHECKING
+
+from classiq.interface.generator.expressions.expression import Expression
+from classiq.interface.generator.functions.port_declaration import (
+    PortDeclarationDirection,
+)
+from classiq.interface.model.quantum_statement import QuantumOperation
+
+from classiq.model_expansions.quantum_operations.emitter import Emitter
+
+if TYPE_CHECKING:
+    from classiq.model_expansions.interpreters.base_interpreter import BaseInterpreter
+
+
+class ExpressionEvaluator(Emitter[QuantumOperation]):
+    def __init__(self, interpreter: "BaseInterpreter", expression_name: str) -> None:
+        super().__init__(interpreter)
+        self._expression_name = expression_name
+
+    def emit(self, op: QuantumOperation, /) -> bool:
+        expression = getattr(op, self._expression_name)
+        if not isinstance(expression, Expression) or expression.is_evaluated():
+            return False
+        evaluated_expression = self._evaluate_expression(
+            expression, preserve_bool_ops=True
+        )
+        for symbol in self._get_symbols_in_expression(evaluated_expression):
+            self._capture_handle(symbol.handle, PortDeclarationDirection.Inout)
+        op = op.model_copy(
+            update={self._expression_name: evaluated_expression, "back_ref": op.uuid}
+        )
+        self._interpreter.emit(op)
+        return True

classiq/model_expansions/quantum_operations/handle_evaluator.py

@@ -0,0 +1,28 @@
+from typing import TYPE_CHECKING
+
+from classiq.interface.model.handle_binding import HandleBinding
+from classiq.interface.model.quantum_statement import QuantumOperation
+
+from classiq.model_expansions.quantum_operations.emitter import Emitter
+
+if TYPE_CHECKING:
+    from classiq.model_expansions.interpreters.base_interpreter import BaseInterpreter
+
+
+class HandleEvaluator(Emitter[QuantumOperation]):
+    def __init__(self, interpreter: "BaseInterpreter", handle_name: str) -> None:
+        super().__init__(interpreter)
+        self._handle_name = handle_name
+
+    def emit(self, op: QuantumOperation, /) -> bool:
+        handle = getattr(op, self._handle_name)
+        if not isinstance(handle, HandleBinding):
+            return False
+        evaluated_handle = self._interpreter.evaluate(handle).value.handle
+        if handle == evaluated_handle:
+            return False
+        op = op.model_copy(
+            update={self._handle_name: evaluated_handle, "back_ref": op.uuid}
+        )
+        self._interpreter.emit(op)
+        return True

classiq/model_expansions/quantum_operations/quantum_function_call.py

@@ -28,16 +28,17 @@ class QuantumFunctionCallEmitter(CallEmitter[QuantumFunctionCall]):
         super().__init__(interpreter)
         self._model = self._interpreter._model
 
-    def emit(self, call: QuantumFunctionCall, /) -> None:
+    def emit(self, call: QuantumFunctionCall, /) -> bool:
         if call.function == "allocate":  # FIXME: Remove compatibility (CAD-25935)
             self._allocate_compatibility(call)
-            return
+            return True
         function = self._interpreter.evaluate(call.function).as_type(FunctionClosure)
         args = call.positional_args
         with ErrorManager().call(function.name):
             self._emit_quantum_function_call(
                 function, args, self._debug_info.get(call.uuid)
             )
+        return True
 
     def _allocate_compatibility(self, call: QuantumFunctionCall) -> None:
         if len(call.positional_args) != 2:

classiq/model_expansions/quantum_operations/repeat.py

@@ -1,3 +1,4 @@
+from classiq.interface.debug_info.debug_info import new_function_debug_info_by_node
 from classiq.interface.exceptions import ClassiqExpansionError
 from classiq.interface.generator.expressions.expression import Expression
 from classiq.interface.generator.functions.builtins.internal_operators import (
@@ -16,7 +17,7 @@ from classiq.qmod.quantum_function import GenerativeQFunc
 
 
 class RepeatEmitter(CallEmitter[Repeat]):
-    def emit(self, repeat: Repeat, /) -> None:
+    def emit(self, repeat: Repeat, /) -> bool:
         count = self._interpreter.evaluate(repeat.count).as_type(int)
         if count < 0:
             raise ClassiqExpansionError(
@@ -25,6 +26,7 @@ class RepeatEmitter(CallEmitter[Repeat]):
         op_name = self._counted_name_allocator.allocate(REPEAT_OPERATOR_NAME)
         for i in range(count):
             self._emit_iteration(repeat, i, op_name)
+        return True
 
     def _emit_iteration(self, repeat: Repeat, i: int, op_name: str) -> None:
         closure_constructor: type[FunctionClosure]
@@ -50,7 +52,11 @@ class RepeatEmitter(CallEmitter[Repeat]):
             ],
             body=repeat.body,
             scope=Scope(parent=self._current_scope),
-            is_lambda=True,
+            lambda_external_vars=self._builder.current_block.captured_vars,
             **extra_args,
         )
-        self._emit_quantum_function_call(iteration_function, [Expression(expr=str(i))])
+        self._emit_quantum_function_call(
+            iteration_function,
+            [Expression(expr=str(i))],
+            new_function_debug_info_by_node(repeat),
+        )

classiq/model_expansions/quantum_operations/variable_decleration.py

@@ -1,3 +1,4 @@
+from classiq.interface.exceptions import ClassiqExpansionError
 from classiq.interface.model.handle_binding import HandleBinding
 from classiq.interface.model.variable_declaration_statement import (
     VariableDeclarationStatement,
@@ -11,9 +12,15 @@ from classiq.model_expansions.scope import Evaluated, QuantumSymbol
 
 
 class VariableDeclarationStatementEmitter(Emitter[VariableDeclarationStatement]):
-    def emit(self, variable_declaration: VariableDeclarationStatement, /) -> None:
-        var_decl = variable_declaration.model_copy()
+    def emit(self, variable_declaration: VariableDeclarationStatement, /) -> bool:
+        var_decl = variable_declaration.model_copy(
+            update=dict(back_ref=variable_declaration.uuid)
+        )
         var_decl.quantum_type = variable_declaration.quantum_type.model_copy()
+        if variable_declaration.name in self._current_scope:
+            raise ClassiqExpansionError(
+                f"Variable {variable_declaration.name!r} is already defined"
+            )
         self._current_scope[variable_declaration.name] = Evaluated(
             value=QuantumSymbol(
                 handle=HandleBinding(name=var_decl.name),
@@ -25,4 +32,8 @@ class VariableDeclarationStatementEmitter(Emitter[VariableDeclarationStatement])
             ),
             defining_function=self._builder.current_function,
         )
+        self._builder.current_block.captured_vars.init_var(
+            var_decl.name, self._builder.current_function
+        )
         self.emit_statement(var_decl)
+        return True

classiq/model_expansions/sympy_conversion/expression_to_sympy.py

@@ -174,6 +174,6 @@ class ExpressionSympyTranslator(ast.NodeTransformer):
     def visit_Attribute(self, node: ast.Attribute) -> ast.Call:
         return ast.Call(
             func=ast.Name("get_field"),
-            args=[node.value, ast.Constant(value=node.attr)],
+            args=[self.visit(node.value), ast.Constant(value=node.attr)],
             keywords=[],
         )

classiq/open_library/functions/__init__.py

@@ -10,7 +10,7 @@ from .hea import *
 from .linear_pauli_rotation import *
 from .linear_pauli_rotation import _single_pauli
 from .modular_exponentiation import *
-from .modular_exponentiation import _check_msb, _ctrl_x
+from .modular_exponentiation import _check_msb
 from .qaoa_penalty import *
 from .qft_functions import *
 from .qpe import *
@@ -72,7 +72,6 @@ OPEN_LIBRARY_FUNCTIONS = [
     qst_type2,
     modular_increment,
     qft,
-    _ctrl_x,
     _prepare_uniform_trimmed_state_step,
     _qct_d_operator,
     _qct_pi_operator,

classiq/open_library/functions/amplitude_amplification.py

@@ -4,7 +4,7 @@ from classiq.qmod.builtins.operations import (
     allocate,
     bind,
     control,
-    repeat,
+    power,
     within_apply,
 )
 from classiq.qmod.cparam import CInt, CReal
@@ -24,7 +24,7 @@ def amplitude_amplification(
     """
     [Qmod Classiq-library function]
 
-    Applies the Amplitude Amplification algorithm (QAE); Prepares a state using the given `space_transform` function, and applies `reps` repetititions
+    Applies the Amplitude Amplification algorithm; Prepares a state using the given `space_transform` function, and applies `reps` repetititions
     of the grover operator, using the given `oracle` functions which marks the "good" states.
 
     Args:
@@ -34,7 +34,10 @@ def amplitude_amplification(
         packed_vars: The variable that holds the state to be amplified. Assumed to be in the zero state at the beginning of the algorithm.
     """
     space_transform(packed_qvars)
-    repeat(reps, lambda index: grover_operator(oracle, space_transform, packed_qvars))
+    power(
+        reps,
+        lambda: grover_operator(oracle, space_transform, packed_qvars),
+    )
 
 
 @qfunc
@@ -47,18 +50,18 @@ def exact_amplitude_amplification(
     """
     [Qmod Classiq-library function]
 
-    Applies an exact version of the Amplitude Amplification algorithm (QAE), assuming knowledge of the amplitude of the marked state.
+    Applies an exact version of the Amplitude Amplification algorithm, assuming knowledge of the amplitude of the marked state.
     The function should be applied on the zero state, and it takes care for preparing the initial state before amplification using the `space_transform`.
 
     Based on the algorithm in [Quantum state preparation without coherent arithmetic](https://arxiv.org/abs/2210.14892).
 
-    Assuming the `space_transform` creates a state $|\\psi\rangle = a|\\psi_good\rangle + \\sqrt(1-a)|\\psi_bad\rangle$, given `a` as the `amplitude`
-    argument, the function will load exactly the state $|\\psi_good\rangle$.
+    Assuming the `space_transform` creates a state $|\\psi\\rangle = a|\\psi_{good}\\rangle + \\sqrt{1-a}|\\psi_{bad}\\rangle$, given `a` as the `amplitude`
+    argument, the function will load exactly the state $|\\psi_{good}\\rangle$.
 
-    Note: if the `amplitude` argument is not exact, the resulting state will not be exactly $|\\psi_good\rangle$, and there will be additional internal auxilliary of the function that is not released correctly.
+    Note: if the `amplitude` argument is not exact, the resulting state will not be exactly $|\\psi_{good}\\rangle$, and there will be additional internal auxilliary of the function that is not released correctly.
 
     Args:
-        amplitude: The amplitude of the state $|\\psi_good\rangle$ with regards to the initial state prepared by.
+        amplitude: The amplitude of the state $|\\psi_{good}\\rangle$ with regards to the initial state prepared by `space_transform`.
         oracle: The oracle operator that marks the "good" states. This operator should flip the sign of the amplitude of the "good" state.
         space_transform: The space transform operator (which is known also the state preparation operator). First applied to prepare the state before the amplification, then used inside the Grover operator.
         packed_vars: The variable that holds the state to be amplified. Assumed to be in the zero state at the beginning of the algorithm.
@@ -68,7 +71,7 @@ def exact_amplitude_amplification(
     theta = pi / (4 * k + 2)
     rot_phase = 2 * acos(sin(theta) / amplitude)
 
-    extended_qvars: QArray = QArray("extended_qvars")
+    extended_qvars: QArray = QArray()
     within_apply(
         lambda: [  # type:ignore[arg-type]
             allocate(aux),