Add basis rotation instructions

README.md

@@ -42,11 +42,16 @@ print(program.json(indent=2))
       "type": "cnot"
     }
   ],
-  "results": null
+  "results": null,
+  "basis_rotation_instructions": null
 }
 """
 
-program = Program(instructions=[H(target=0), CNot(control=0, target=1)], results=[Expectation(targets=[1], observable=['x'])])
+program = Program(
+    instructions=[H(target=0), CNot(control=0, target=1)],
+    results=[Expectation(targets=[0], observable=['x'])],
+    basis_rotation_instructions=[H(target=0)]
+)
 print(program.json(indent=2))
 
 """
@@ -68,10 +73,16 @@ print(program.json(indent=2))
         "x"
       ],
       "targets": [
-        1
+        0
       ],
       "type": "expectation"
     }
+  ],
+  "basis_rotation_instructions": [
+    {
+      "target": 0,
+      "type": "h"
+    }
   ]
 }
 """
@@ -112,10 +123,16 @@ jaqcd_string = """
         "x"
       ],
       "targets": [
-        1
+        0
       ],
       "type": "expectation"
     }
+  ],
+  "basis_rotation_instructions": [
+    {
+      "target": 0,
+      "type": "h"
+    }
   ]
 }
 """
@@ -124,7 +141,7 @@ program = Program.parse_raw(jaqcd_string)
 print(program)
 
 """
-instructions=[H(target=0, type=<Type.h: 'h'>), CNot(control=0, target=1, type=<Type.cnot: 'cnot'>)] results=[Expectation(observable=['x'], targets=[1], type=<Type.expectation: 'expectation'>)]
+instructions=[H(target=0, type=<Type.h: 'h'>), CNot(control=0, target=1, type=<Type.cnot: 'cnot'>)] results=[Expectation(observable=['x'], targets=[0], type=<Type.expectation: 'expectation'>)] basis_rotation_instructions=[H(target=0, type=<Type.h: 'h'>)]
 """
 
 annealing_string = """

setup.py

@@ -32,7 +32,7 @@
             "pytest-cov",
             "pytest-rerunfailures",
             "pytest-xdist",
-            "sphinx",
+            "sphinx < 3.0.0",
             "sphinx-rtd-theme",
             "tox",
         ]

src/braket/ir/jaqcd/program.py

@@ -57,54 +57,99 @@
 )
 from pydantic import BaseModel
 
+GateInstructions = Union[
+    CCNot,
+    CNot,
+    CPhaseShift,
+    CPhaseShift00,
+    CPhaseShift01,
+    CPhaseShift10,
+    CSwap,
+    CY,
+    CZ,
+    H,
+    I,
+    ISwap,
+    PhaseShift,
+    PSwap,
+    Rx,
+    Ry,
+    Rz,
+    S,
+    Swap,
+    Si,
+    T,
+    Ti,
+    Unitary,
+    V,
+    Vi,
+    X,
+    XX,
+    XY,
+    Y,
+    YY,
+    Z,
+    ZZ,
+]
+
 
 class Program(BaseModel):
     """
     Root object of the JsonAwsQuantumCircuitDescription IR.
 
+
+
     Attributes:
         - instructions: List of instructions.
+        - basis_rotation_instructions: List of instructions for rotation to desired measurement
+            bases
+        - results: List of requested results
 
     Examples:
         >>> Program(instructions=[H(target=0), Rz(angle=0.15, target=1)])
+        >>> Program(instructions=[H(target=0), CNot(control=0, target=1)],
+        ...     results=[Expectation(targets=[0], observable=['x'])],
+        ...     basis_rotation_instructions=[H(target=0)])
+
+
+    Note:
+        The type `GateInstructions` includes the following instructions:
+        CCNot,
+        CNot,
+        CPhaseShift,
+        CPhaseShift00,
+        CPhaseShift01,
+        CPhaseShift10,
+        CSwap,
+        CY,
+        CZ,
+        H,
+        I,
+        ISwap,
+        PhaseShift,
+        PSwap,
+        Rx,
+        Ry,
+        Rz,
+        S,
+        Swap,
+        Si,
+        T,
+        Ti,
+        Unitary,
+        V,
+        Vi,
+        X,
+        XX,
+        XY,
+        Y,
+        YY,
+        Z,
+        ZZ
     """
 
-    instructions: List[
-        Union[
-            CCNot,
-            CNot,
-            CPhaseShift,
-            CPhaseShift00,
-            CPhaseShift01,
-            CPhaseShift10,
-            CSwap,
-            CY,
-            CZ,
-            H,
-            I,
-            ISwap,
-            PhaseShift,
-            PSwap,
-            Rx,
-            Ry,
-            Rz,
-            S,
-            Swap,
-            Si,
-            T,
-            Ti,
-            Unitary,
-            V,
-            Vi,
-            X,
-            XX,
-            XY,
-            Y,
-            YY,
-            Z,
-            ZZ,
-        ]
-    ]
+    instructions: List[GateInstructions]
     results: Optional[
         List[Union[Amplitude, Expectation, Probability, Sample, StateVector, Variance]]
     ]
+    basis_rotation_instructions: Optional[List[GateInstructions]]

test/braket/ir/jaqcd/test_program.py

@@ -12,7 +12,7 @@
 # language governing permissions and limitations under the License.
 
 import pytest
-from braket.ir.jaqcd import CNot, Expectation, Program
+from braket.ir.jaqcd import CNot, Expectation, H, Program
 from pydantic import ValidationError
 
 
@@ -48,3 +48,19 @@ def test_instruction_with_results():
         instructions=[CNot(control=0, target=1)],
         results=[Expectation(targets=[1], observable=["x"])],
     )
+
+
+@pytest.mark.xfail(raises=ValidationError)
+def test_partial_non_rotation_basis_instruction():
+    Program(
+        instructions=[CNot(control=0, target=1)],
+        basis_rotation_instructions=[Expectation(targets=[1], observable=["x"]), H(target=1)],
+    )
+
+
+def test_no_rotation_basis_instruction():
+    Program(instructions=[CNot(control=0, target=1)],)
+
+
+def test_rotation_basis_instruction():
+    Program(instructions=[CNot(control=0, target=1)], basis_rotation_instructions=[H(target=1)])