fix: update numpy random function to use a generator (#274)

amazon-braket · Nov 22, 2024 · fd1ed94 · fd1ed94
1 parent a9808f3
commit fd1ed94
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 2 deletions.
diff --git a/src/braket/default_simulator/density_matrix_simulation.py b/src/braket/default_simulator/density_matrix_simulation.py
@@ -108,7 +108,8 @@ def _apply_operations(
         return np.reshape(dm_tensor, (2**qubit_count, 2**qubit_count))
 
     def retrieve_samples(self) -> list[int]:
-        return np.random.choice(
+        rng_generator = np.random.default_rng()
+        return rng_generator.choice(
             self._density_matrix.shape[0], p=self.probabilities, size=self._shots
         )
 

diff --git a/src/braket/default_simulator/state_vector_simulation.py b/src/braket/default_simulator/state_vector_simulation.py
@@ -107,7 +107,8 @@ def _apply_operations(
         return np.reshape(final, 2**qubit_count)
 
     def retrieve_samples(self) -> list[int]:
-        return np.random.choice(len(self._state_vector), p=self.probabilities, size=self._shots)
+        rng_generator = np.random.default_rng()
+        return rng_generator.choice(len(self._state_vector), p=self.probabilities, size=self._shots)
 
     @property
     def state_vector(self) -> np.ndarray: