fix: apply pow modifier to cnot and swap gates

src/braket/default_simulator/gate_operations.py

@@ -208,8 +208,8 @@ def _base_matrix(self) -> np.ndarray:
         return CX._matrix
 
     @property
-    def gate_type(self) -> str:
-        return "cx"
+    def gate_type(self) -> str | None:
+        return "cx" if self._power % 2 == 1 else None
 
 
 @_from_braket_instruction.register(braket_instruction.CNot)
@@ -694,8 +694,8 @@ def _base_matrix(self) -> np.ndarray:
         return Swap._matrix
 
     @property
-    def gate_type(self) -> str:
-        return "swap"
+    def gate_type(self) -> str | None:
+        return "swap" if self._power % 2 == 1 else None
 
 
 @_from_braket_instruction.register(braket_instruction.Swap)

test/unit_tests/braket/default_simulator/test_density_matrix_simulation.py

@@ -18,6 +18,7 @@
 
 from braket.default_simulator import gate_operations, noise_operations, observables
 from braket.default_simulator.density_matrix_simulation import DensityMatrixSimulation
+from braket.default_simulator.state_vector_simulation import StateVectorSimulation
 
 sx = np.array([[0, 1], [1, 0]], dtype=complex)
 si = np.array([[1, 0], [0, 1]], dtype=complex)
@@ -520,3 +521,52 @@ def test_toffoli_gate():
         prob_7 = counter[7] / total_shots
         assert abs(prob_2 - 0.5) < 0.1
         assert abs(prob_7 - 0.5) < 0.1
+
+
+_POWERS = [0, 1, 2, 3, -1, -2, 0.5, -0.5, 1.5]
+
+_POWERED_GATES = [
+    (gate_operations.Hadamard, 1, ()),
+    (gate_operations.PauliX, 1, ()),
+    (gate_operations.PauliZ, 1, ()),
+    (gate_operations.S, 1, ()),
+    (gate_operations.T, 1, ()),
+    (gate_operations.RotX, 1, (0.37,)),
+    (gate_operations.RotZ, 1, (0.37,)),
+    (gate_operations.PhaseShift, 1, (0.37,)),
+    (gate_operations.CX, 2, ()),
+    (gate_operations.CY, 2, ()),
+    (gate_operations.CZ, 2, ()),
+    (gate_operations.Swap, 2, ()),
+    (gate_operations.ISwap, 2, ()),
+    (gate_operations.ECR, 2, ()),
+    (gate_operations.CPhaseShift, 2, (0.37,)),
+    (gate_operations.PSwap, 2, (0.37,)),
+    (gate_operations.XX, 2, (0.37,)),
+    (gate_operations.ZZ, 2, (0.37,)),
+]
+
+
+def _preparation_operations(qubit_count):
+    return [gate_operations.RotY([q], 0.3 + 0.4 * q) for q in range(qubit_count)] + [
+        gate_operations.RotZ([q], 0.2 + 0.3 * q) for q in range(qubit_count)
+    ]
+
+
+def _assert_matches_state_vector(operations, qubit_count):
+    simulation = DensityMatrixSimulation(qubit_count, 0)
+    simulation.evolve(operations)
+    reference = StateVectorSimulation(qubit_count, 0, 1)
+    reference.evolve(operations)
+    state_vector = reference.state_vector
+    assert np.allclose(simulation.density_matrix, np.outer(state_vector, state_vector.conj()))
+
+
+@pytest.mark.parametrize("qubit_count, targets", [(3, (2, 0)), (5, (4, 1))])
+@pytest.mark.parametrize("power", _POWERS)
+@pytest.mark.parametrize("gate_class, num_targets, args", _POWERED_GATES)
+def test_powered_gate(gate_class, num_targets, args, power, qubit_count, targets):
+    operations = _preparation_operations(qubit_count) + [
+        gate_class(targets[:num_targets], *args, power=power)
+    ]
+    _assert_matches_state_vector(operations, qubit_count)

test/unit_tests/braket/default_simulator/test_density_matrix_simulator.py

@@ -20,6 +20,7 @@
 
 import numpy as np
 import pytest
+from scipy.linalg import fractional_matrix_power
 from braket.device_schema.simulators import (
     GateModelSimulatorDeviceCapabilities,
     GateModelSimulatorDeviceParameters,
@@ -1088,3 +1089,42 @@ def test_remove_verbatim_box_with_nested_braces():
     assert "box{" not in result
     assert "if (b == 1) {" in result
     assert "x $0;" in result
+
+
+_TWO_QUBIT_BASE_MATRICES = {
+    "cnot": np.array([[1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 0, 1], [0, 0, 1, 0]], dtype=complex),
+    "swap": np.array([[1, 0, 0, 0], [0, 0, 1, 0], [0, 1, 0, 0], [0, 0, 0, 1]], dtype=complex),
+    "cz": np.diag([1, 1, 1, -1]).astype(complex),
+    "iswap": np.array([[1, 0, 0, 0], [0, 0, 1j, 0], [0, 1j, 0, 0], [0, 0, 0, 1]], dtype=complex),
+}
+
+
+def _openqasm_matrix(matrix):
+    return (
+        "["
+        + ", ".join("[" + ", ".join(f"{z.real}{z.imag:+}im" for z in row) + "]" for row in matrix)
+        + "]"
+    )
+
+
+@pytest.mark.parametrize("power", [0, 2, 3, -1, -2, 0.5, -0.5, 1.5])
+@pytest.mark.parametrize("gate", list(_TWO_QUBIT_BASE_MATRICES))
+def test_pow_two_qubit_gate_matches_unitary(gate, power):
+    preparation = "qubit[2] q;\nry(0.7) q[0];\nry(0.3) q[1];\n"
+    powered = f"OPENQASM 3.0;\n{preparation}pow({power}) @ {gate} q[0], q[1];\n"
+    matrix = fractional_matrix_power(_TWO_QUBIT_BASE_MATRICES[gate], power)
+    unitary = (
+        f"OPENQASM 3.0;\n{preparation}"
+        f"#pragma braket unitary({_openqasm_matrix(matrix)}) q[0], q[1]\n"
+    )
+    result_type = "#pragma braket result density_matrix\n"
+    results = [
+        np.asarray(
+            DensityMatrixSimulator()
+            .run(OpenQASMProgram(source=source + result_type), shots=0)
+            .resultTypes[0]
+            .value
+        )
+        for source in (powered, unitary)
+    ]
+    assert np.allclose(*results)

test/unit_tests/braket/default_simulator/test_gate_operations.py

@@ -88,6 +88,29 @@ def test_gate_operation(ir_instruction, targets, operation_type):
     check_unitary(operation_instance.matrix)
 
 
+@pytest.mark.parametrize(
+    "operation_type, gate_type", [(gate_operations.CX, "cx"), (gate_operations.Swap, "swap")]
+)
+@pytest.mark.parametrize(
+    "power, uses_gate_type",
+    [
+        (1, True),
+        (3, True),
+        (-1, True),
+        (1.0, True),
+        (0, False),
+        (2, False),
+        (-2, False),
+        (0.5, False),
+    ],
+)
+def test_matrix_free_gate_type_respects_power(operation_type, gate_type, power, uses_gate_type):
+    operation = operation_type([0, 1], power=power)
+    assert operation.gate_type == (gate_type if uses_gate_type else None)
+    if uses_gate_type:
+        np.testing.assert_allclose(operation.matrix, operation._base_matrix)
+
+
 # ---------------------------------------------------------------------------
 # Measure class tests
 # ---------------------------------------------------------------------------

test/unit_tests/braket/default_simulator/test_state_vector_simulation.py

@@ -17,6 +17,7 @@
 import pytest
 
 from braket.default_simulator import gate_operations, observables
+from braket.default_simulator.linalg_utils import controlled_matrix
 from braket.default_simulator.state_vector_simulation import StateVectorSimulation
 
 evolve_testdata = [
@@ -544,3 +545,85 @@ def test_ghz_tree_retrieve_samples(batch_size):
     assert 0.4 < counter[0] / (counter[0] + counter[960]) < 0.6
     assert 0.4 < counter[960] / (counter[0] + counter[960]) < 0.6
     assert counter[0] + counter[960] == 10000
+
+
+_POWERS = [0, 1, 2, 3, -1, -2, 0.5, -0.5, 1.5]
+
+_POWERED_GATES = [
+    (gate_operations.Hadamard, 1, ()),
+    (gate_operations.PauliX, 1, ()),
+    (gate_operations.PauliZ, 1, ()),
+    (gate_operations.S, 1, ()),
+    (gate_operations.T, 1, ()),
+    (gate_operations.RotX, 1, (0.37,)),
+    (gate_operations.RotZ, 1, (0.37,)),
+    (gate_operations.PhaseShift, 1, (0.37,)),
+    (gate_operations.CX, 2, ()),
+    (gate_operations.CY, 2, ()),
+    (gate_operations.CZ, 2, ()),
+    (gate_operations.Swap, 2, ()),
+    (gate_operations.ISwap, 2, ()),
+    (gate_operations.ECR, 2, ()),
+    (gate_operations.CPhaseShift, 2, (0.37,)),
+    (gate_operations.PSwap, 2, (0.37,)),
+    (gate_operations.XX, 2, (0.37,)),
+    (gate_operations.ZZ, 2, (0.37,)),
+]
+
+
+def _preparation_operations(qubit_count):
+    return [gate_operations.RotY([q], 0.3 + 0.4 * q) for q in range(qubit_count)] + [
+        gate_operations.RotZ([q], 0.2 + 0.3 * q) for q in range(qubit_count)
+    ]
+
+
+def _apply_reference(state_vector, operation, qubit_count):
+    matrix = controlled_matrix(operation.matrix, operation.control_state)
+    targets = operation.targets
+    num_targets = len(targets)
+    product = np.tensordot(
+        matrix.reshape([2] * 2 * num_targets),
+        state_vector.reshape([2] * qubit_count),
+        axes=(range(num_targets, 2 * num_targets), targets),
+    )
+    return np.moveaxis(product, range(num_targets), targets).reshape(-1)
+
+
+def _reference_state_vector(operations, qubit_count):
+    state_vector = np.zeros(2**qubit_count, dtype=complex)
+    state_vector[0] = 1
+    for operation in operations:
+        state_vector = _apply_reference(state_vector, operation, qubit_count)
+    return state_vector
+
+
+@pytest.mark.parametrize("qubit_count, targets", [(3, (2, 0)), (11, (9, 3))])
+@pytest.mark.parametrize("power", _POWERS)
+@pytest.mark.parametrize("gate_class, num_targets, args", _POWERED_GATES)
+def test_powered_gate(gate_class, num_targets, args, power, qubit_count, targets):
+    operations = _preparation_operations(qubit_count) + [
+        gate_class(targets[:num_targets], *args, power=power)
+    ]
+    simulation = StateVectorSimulation(qubit_count, 0, 1)
+    simulation.evolve(operations)
+    assert np.allclose(simulation.state_vector, _reference_state_vector(operations, qubit_count))
+
+
+@pytest.mark.parametrize(
+    "qubit_count, targets, control_state",
+    [
+        (4, (1, 3, 0), (1,)),
+        (4, (1, 3, 0), (0,)),
+        (4, (1, 3, 0, 2), (1, 1)),
+        (4, (3, 1, 2, 0), (0, 0)),
+    ],
+)
+@pytest.mark.parametrize("power", [0, 2, 0.5, -1])
+@pytest.mark.parametrize("gate_class", [gate_operations.CX, gate_operations.Swap])
+def test_controlled_powered_gate(gate_class, power, qubit_count, targets, control_state):
+    operations = _preparation_operations(qubit_count) + [
+        gate_class(targets, ctrl_modifiers=control_state, power=power)
+    ]
+    simulation = StateVectorSimulation(qubit_count, 0, 1)
+    simulation.evolve(operations)
+    assert np.allclose(simulation.state_vector, _reference_state_vector(operations, qubit_count))

test/unit_tests/braket/default_simulator/test_state_vector_simulator.py

@@ -20,6 +20,7 @@
 
 import numpy as np
 import pytest
+from scipy.linalg import fractional_matrix_power
 from braket.device_schema.simulators import (
     GateModelSimulatorDeviceCapabilities,
     GateModelSimulatorDeviceParameters,
@@ -1707,3 +1708,42 @@ def test_program_set_validation():
 
     result = simulator.run(multi_program_set, shots=10)
     assert result is not None
+
+
+_TWO_QUBIT_BASE_MATRICES = {
+    "cnot": np.array([[1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 0, 1], [0, 0, 1, 0]], dtype=complex),
+    "swap": np.array([[1, 0, 0, 0], [0, 0, 1, 0], [0, 1, 0, 0], [0, 0, 0, 1]], dtype=complex),
+    "cz": np.diag([1, 1, 1, -1]).astype(complex),
+    "iswap": np.array([[1, 0, 0, 0], [0, 0, 1j, 0], [0, 1j, 0, 0], [0, 0, 0, 1]], dtype=complex),
+}
+
+
+def _openqasm_matrix(matrix):
+    return (
+        "["
+        + ", ".join("[" + ", ".join(f"{z.real}{z.imag:+}im" for z in row) + "]" for row in matrix)
+        + "]"
+    )
+
+
+@pytest.mark.parametrize("power", [0, 2, 3, -1, -2, 0.5, -0.5, 1.5])
+@pytest.mark.parametrize("gate", list(_TWO_QUBIT_BASE_MATRICES))
+def test_pow_two_qubit_gate_matches_unitary(gate, power):
+    preparation = "qubit[2] q;\nry(0.7) q[0];\nry(0.3) q[1];\n"
+    powered = f"OPENQASM 3.0;\n{preparation}pow({power}) @ {gate} q[0], q[1];\n"
+    matrix = fractional_matrix_power(_TWO_QUBIT_BASE_MATRICES[gate], power)
+    unitary = (
+        f"OPENQASM 3.0;\n{preparation}"
+        f"#pragma braket unitary({_openqasm_matrix(matrix)}) q[0], q[1]\n"
+    )
+    result_type = "#pragma braket result state_vector\n"
+    results = [
+        np.asarray(
+            StateVectorSimulator()
+            .run(OpenQASMProgram(source=source + result_type), shots=0)
+            .resultTypes[0]
+            .value
+        )
+        for source in (powered, unitary)
+    ]
+    assert np.allclose(*results)