fix: correct parameter types in sampler methods (#65)

amazon-braket · Feb 9, 2022 · 35522f9 · 35522f9
1 parent 2f46afd
commit 35522f9
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 8 deletions.
diff --git a/src/braket/ocean_plugin/braket_dwave_sampler.py b/src/braket/ocean_plugin/braket_dwave_sampler.py
@@ -110,7 +110,7 @@ def parameters(self) -> FrozenDict[str, List]:
         )
 
     def sample_ising(
-        self, h: Union[Dict[int, int], List[int]], J: Dict[int, int], **kwargs
+        self, h: Union[Dict[int, float], List[float]], J: Dict[Tuple[int, int], float], **kwargs
     ) -> SampleSet:
         """
         Sample from the specified Ising model.
@@ -162,7 +162,7 @@ def sample_ising(
         return super().sample_ising(h, J, **kwargs)
 
     def sample_ising_quantum_task(
-        self, h: Union[Dict[int, int], List[int]], J: Dict[Tuple[int, int], float], **kwargs
+        self, h: Union[Dict[int, float], List[float]], J: Dict[Tuple[int, int], float], **kwargs
     ) -> QuantumTask:
         """
         Sample from the specified Ising model and return a `QuantumTask`. This has the same inputs
@@ -217,7 +217,7 @@ def sample_ising_quantum_task(
         """
         return super().sample_ising_quantum_task(h, J, **kwargs)
 
-    def sample_qubo(self, Q: Dict[Tuple[int, int], int], **kwargs) -> SampleSet:
+    def sample_qubo(self, Q: Dict[Tuple[int, int], float], **kwargs) -> SampleSet:
         """
         Sample from the specified QUBO.
 
@@ -260,7 +260,7 @@ def sample_qubo(self, Q: Dict[Tuple[int, int], int], **kwargs) -> SampleSet:
         """
         return super().sample_qubo(Q, **kwargs)
 
-    def sample_qubo_quantum_task(self, Q: Dict[Tuple[int, int], int], **kwargs) -> QuantumTask:
+    def sample_qubo_quantum_task(self, Q: Dict[Tuple[int, int], float], **kwargs) -> QuantumTask:
         """
         Sample from the specified QUBO and return a `QuantumTask`. This has the same inputs
         as `BraketDWaveSampler.sample_qubo`.

diff --git a/src/braket/ocean_plugin/braket_sampler.py b/src/braket/ocean_plugin/braket_sampler.py
@@ -147,7 +147,7 @@ def _access_optimized_edgelist(self) -> FrozenDict[int, FrozenSet[int]]:
         return FrozenDict(edges)
 
     def sample_ising(
-        self, h: Union[Dict[int, int], List[int]], J: Dict[int, int], **kwargs
+        self, h: Union[Dict[int, float], List[float]], J: Dict[Tuple[int, int], float], **kwargs
     ) -> SampleSet:
         """
         Sample from the specified Ising model.
@@ -210,7 +210,7 @@ def sample_ising(
         return BraketSampler.get_task_sample_set(aws_task, variables)
 
     def sample_ising_quantum_task(
-        self, h: Union[Dict[int, int], List[int]], J: Dict[Tuple[int, int], float], **kwargs
+        self, h: Union[Dict[int, float], List[float]], J: Dict[Tuple[int, int], float], **kwargs
     ) -> QuantumTask:
         """
         Sample from the specified Ising model and return a `QuantumTask`. This has the same inputs
@@ -285,7 +285,7 @@ def sample_ising_quantum_task(
             **solver_kwargs,
         )
 
-    def sample_qubo(self, Q: Dict[Tuple[int, int], int], **kwargs) -> SampleSet:
+    def sample_qubo(self, Q: Dict[Tuple[int, int], float], **kwargs) -> SampleSet:
         """
         Sample from the specified QUBO.
 
@@ -335,7 +335,7 @@ def sample_qubo(self, Q: Dict[Tuple[int, int], int], **kwargs) -> SampleSet:
         variables = set().union(*Q)
         return BraketSampler.get_task_sample_set(aws_task, variables)
 
-    def sample_qubo_quantum_task(self, Q: Dict[Tuple[int, int], int], **kwargs) -> QuantumTask:
+    def sample_qubo_quantum_task(self, Q: Dict[Tuple[int, int], float], **kwargs) -> QuantumTask:
         """
         Sample from the specified QUBO and return a `QuantumTask`. This has the same inputs
         as `BraketSampler.sample_qubo`