Add header and versioning to IR

README.md

@@ -44,6 +44,10 @@ print(program.json(indent=2))
 
 """
 {
+  "braketSchemaHeader": {
+    "name": "braket.ir.jaqcd.program",
+    "version": "1"
+  },
   "instructions": [
     {
       "target": 0,
@@ -56,7 +60,7 @@ print(program.json(indent=2))
     }
   ],
   "results": null,
-  "basis_rotation_instructions": null
+  "basis_rotation_instructions": null,
 }
 """
 
@@ -69,6 +73,10 @@ print(program.json(indent=2))
 
 """
 {
+  "braketSchemaHeader": {
+    "name": "braket.ir.jaqcd.program",
+    "version": "1"
+  },
   "instructions": [
     {
       "target": 0,
@@ -105,6 +113,10 @@ print(problem.json(indent=2))
 
 """
 {
+  "braketSchemaHeader": {
+    "name": "braket.ir.annealing.problem",
+    "version": "1"
+  },
   "type": "QUBO",
   "linear": {0: 0.3, 4: -0.3},
   "quadratic": {"0,5": 0.667}
@@ -154,7 +166,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=[0], type=<Type.expectation: 'expectation'>)] basis_rotation_instructions=[H(target=0, type=<Type.h: 'h'>)]
+braketSchemaHeader=BraketSchemaHeader(name='braket.ir.jaqcd.program', version='1') 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 = """
@@ -169,7 +181,7 @@ problem = Problem.parse_raw(annealing_string)
 print(problem)
 
 """
-type=<ProblemType.QUBO: 'QUBO'>, linear={0: 0.3, 4: -0.3}, quadratic={'0,5': 0.667}
+braketSchemaHeader=BraketSchemaHeader(name='braket.ir.annealing.problem', version='1') type=<ProblemType.QUBO: 'QUBO'>, linear={0: 0.3, 4: -0.3}, quadratic={'0,5': 0.667}
 """
 
 ```

src/braket/ir/annealing/__init__.py

@@ -11,4 +11,4 @@
 # ANY KIND, either express or implied. See the License for the specific
 # language governing permissions and limitations under the License.
 
-from braket.ir.annealing.problem import Problem, ProblemType  # noqa: F401
+from braket.ir.annealing.problem_v1 import Problem, ProblemType  # noqa: F401

src/braket/ir/annealing/problem.py → src/braket/ir/annealing/problem_v1.py

@@ -14,7 +14,9 @@
 from enum import Enum
 from typing import Dict
 
-from pydantic import BaseModel, conint
+from pydantic import Field, conint
+
+from braket.schema_common import BraketSchemaBase, BraketSchemaHeader
 
 
 class ProblemType(str, Enum):
@@ -28,19 +30,23 @@ class ProblemType(str, Enum):
     ISING = "ISING"
 
 
-class Problem(BaseModel):
+class Problem(BraketSchemaBase):
     """ Specifies a quantum annealing problem.
 
     Attributes:
-        - type: The type of problem; can be either "QUBO" or "ISING"
-        - linear: Linear terms of the model.
-        - quadratic: Quadratic terms of the model, keyed on comma-separated
+        braketSchemaHeader (BraketSchemaHeader): Schema header. Users do not need
+            to set this value. Only default is allowed.
+        type (ProblemType): The type of problem; can be either "QUBO" or "ISING"
+        linear (Dict[int, float]): Linear terms of the model.
+        quadratic (Dict[str, float]): Quadratic terms of the model, keyed on comma-separated
             variables as strings
 
     Examples:
         >>> Problem(type=ProblemType.QUBO, linear={0: 0.3, 4: -0.3}, quadratic={"0,5": 0.667})
     """
 
+    _PROBLEM_HEADER = BraketSchemaHeader(name="braket.ir.annealing.problem", version="1")
+    braketSchemaHeader: BraketSchemaHeader = Field(default=_PROBLEM_HEADER, const=_PROBLEM_HEADER)
     type: ProblemType
     linear: Dict[conint(ge=0), float]
     quadratic: Dict[str, float]

src/braket/ir/jaqcd/__init__.py

@@ -45,7 +45,7 @@
     Y,
     Z,
 )
-from braket.ir.jaqcd.program import Program  # noqa: F401
+from braket.ir.jaqcd.program_v1 import Program  # noqa: F401
 from braket.ir.jaqcd.results import (  # noqa: F401
     Amplitude,
     Expectation,

src/braket/ir/jaqcd/program.py → src/braket/ir/jaqcd/program_v1.py

@@ -13,7 +13,7 @@
 
 from typing import List, Optional, Union
 
-from pydantic import BaseModel
+from pydantic import Field
 
 from braket.ir.jaqcd.instructions import (
     CY,
@@ -57,6 +57,7 @@
     StateVector,
     Variance,
 )
+from braket.schema_common import BraketSchemaBase, BraketSchemaHeader
 
 GateInstructions = Union[
     CCNot,
@@ -94,17 +95,20 @@
 ]
 
 
-class Program(BaseModel):
+class Program(BraketSchemaBase):
     """
     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
+        braketSchemaHeader (BraketSchemaHeader): Schema header. Users do not need
+            to set this value. Only default is allowed.
+        instructions (List[GateInstructions]): List of instructions.
+        basis_rotation_instructions (List[GateInstructions]): List of instructions for
+            rotation to desired measurement bases. Default is None.
+        results (List[Union[Amplitude, Expectation, Probability, Sample, StateVector, Variance]]):
+            List of requested results. Default is None.
 
     Examples:
         >>> Program(instructions=[H(target=0), Rz(angle=0.15, target=1)])
@@ -149,6 +153,8 @@ class Program(BaseModel):
         ZZ
     """
 
+    _PROGRAM_HEADER = BraketSchemaHeader(name="braket.ir.jaqcd.program", version="1")
+    braketSchemaHeader: BraketSchemaHeader = Field(default=_PROGRAM_HEADER, const=_PROGRAM_HEADER)
     instructions: List[GateInstructions]
     results: Optional[
         List[Union[Amplitude, Expectation, Probability, Sample, StateVector, Variance]]

src/braket/task_result/annealing_task_result_v1.py

@@ -38,11 +38,11 @@ class AnnealingTaskResult(BraketSchemaBase):
 
     """
 
-    ANNEALING_TASK_RESULT_HEADER = BraketSchemaHeader(
+    _ANNEALING_TASK_RESULT_HEADER = BraketSchemaHeader(
         name="braket.task_result.annealing_task_result", version="1"
     )
     braketSchemaHeader: BraketSchemaHeader = Field(
-        default=ANNEALING_TASK_RESULT_HEADER, const=ANNEALING_TASK_RESULT_HEADER
+        default=_ANNEALING_TASK_RESULT_HEADER, const=_ANNEALING_TASK_RESULT_HEADER
     )
     solutions: List[conlist(conint(ge=-1, le=3), min_items=1)]
     solutionCounts: Optional[List[conint(ge=0)]]

src/braket/task_result/dwave_metadata_v1.py

@@ -58,11 +58,11 @@ class DwaveMetadata(BraketSchemaBase):
         >>> DwaveMetadata(activeVariables=[0, 3, 4], timing=timing)
     """
 
-    DWAVE_METADATA_HEADER = BraketSchemaHeader(
+    _DWAVE_METADATA_HEADER = BraketSchemaHeader(
         name="braket.task_result.dwave_metadata", version="1"
     )
     braketSchemaHeader: BraketSchemaHeader = Field(
-        default=DWAVE_METADATA_HEADER, const=DWAVE_METADATA_HEADER
+        default=_DWAVE_METADATA_HEADER, const=_DWAVE_METADATA_HEADER
     )
     activeVariables: conlist(conint(ge=0))
     timing: DwaveTiming

src/braket/task_result/gate_model_task_result_v1.py

@@ -55,12 +55,12 @@ class GateModelTaskResult(BraketSchemaBase):
         additionalMetadata (AdditionalMetadata): Additional metadata of the task
     """
 
-    GATE_MODEL_TASK_RESULT_HEADER = BraketSchemaHeader(
+    _GATE_MODEL_TASK_RESULT_HEADER = BraketSchemaHeader(
         name="braket.task_result.gate_model_task_result", version="1"
     )
 
     braketSchemaHeader: BraketSchemaHeader = Field(
-        default=GATE_MODEL_TASK_RESULT_HEADER, const=GATE_MODEL_TASK_RESULT_HEADER
+        default=_GATE_MODEL_TASK_RESULT_HEADER, const=_GATE_MODEL_TASK_RESULT_HEADER
     )
     measurements: Optional[conlist(conlist(conint(ge=0, le=1), min_items=1), min_items=1)]
     measurementProbabilities: Optional[

src/braket/task_result/task_metadata_v1.py

@@ -45,10 +45,10 @@ class TaskMetadata(BraketSchemaBase):
 
     """
 
-    TASK_METADATA_HEADER = BraketSchemaHeader(name="braket.task_result.task_metadata", version="1")
+    _TASK_METADATA_HEADER = BraketSchemaHeader(name="braket.task_result.task_metadata", version="1")
 
     braketSchemaHeader: BraketSchemaHeader = Field(
-        default=TASK_METADATA_HEADER, const=TASK_METADATA_HEADER
+        default=_TASK_METADATA_HEADER, const=_TASK_METADATA_HEADER
     )
     id: constr(min_length=1)
     shots: conint(ge=0)

test/braket/ir/annealing/test_problem.py → test/braket/ir/annealing/test_problem_v1.py

@@ -14,7 +14,7 @@
 import pytest
 from pydantic import ValidationError
 
-from braket.ir.annealing import Problem, ProblemType
+from braket.ir.annealing.problem_v1 import Problem, ProblemType
 
 
 def test_creation():
@@ -27,6 +27,7 @@ def test_creation():
     assert problem.linear == {0: 0.3333, 1: -0.333, 4: -0.333, 5: 0.333}
     assert problem.quadratic == {"0,4": 0.667, "0,5": -1, "1,4": 0.667, "1,5": 0.667}
     assert Problem.parse_raw(problem.json()) == problem
+    assert problem == Problem.parse_raw_schema(problem.json())
 
 
 @pytest.mark.xfail(raises=ValidationError)

test/braket/ir/jaqcd/test_program.py → test/braket/ir/jaqcd/test_program_v1.py

@@ -14,7 +14,8 @@
 import pytest
 from pydantic import ValidationError
 
-from braket.ir.jaqcd import CNot, Expectation, H, Program
+from braket.ir.jaqcd import CNot, Expectation, H
+from braket.ir.jaqcd.program_v1 import Program
 
 
 @pytest.mark.xfail(raises=ValidationError)
@@ -41,7 +42,8 @@ def test_partial_non_result():
 
 
 def test_instruction_no_results():
-    Program(instructions=[CNot(control=0, target=1)])
+    program = Program(instructions=[CNot(control=0, target=1)])
+    assert Program.parse_raw(program.json()) == program
 
 
 def test_instruction_with_results():

test/braket/ir/jaqcd/test_results.py

@@ -10,7 +10,7 @@
     idfn,
 )
 
-from braket.ir.jaqcd.program import Program
+from braket.ir.jaqcd import Program
 from braket.ir.jaqcd.results import (
     Amplitude,
     Expectation,