diff --git a/src/qforte/qkd/srqk.py b/src/qforte/qkd/srqk.py
index 05dc3764..8ac3bcf5 100644
--- a/src/qforte/qkd/srqk.py
+++ b/src/qforte/qkd/srqk.py
@@ -137,8 +137,7 @@ def build_qk_mats_fast(self):
 
         for m in range(self._nstates):
             # Compute U_m = exp(-i m dt H)
-            Um = qforte.Circuit()
-            Um.add(self._Uprep)
+            Um = qforte.Circuit(self._Uprep)
             phase1 = 1.0
 
             if m > 0:
@@ -234,6 +233,9 @@ def matrix_element(self, m, n, use_op=False):
         n : int
             The number of time steps for the Un evolution.
 
+        use_op : bool
+            Should the matrix element be of H (true) or S (false)?
+
         Returns
         -------
         value : complex
@@ -277,11 +279,7 @@ def matrix_element(self, m, n, use_op=False):
             Ub.add(gate)
 
         if not use_op:
-            # TODO (opt): use Uprep
-            cir = qforte.Circuit()
-            for j in range(self._nqb):
-                if self._ref[j] == 1:
-                    cir.add(qforte.gate("X", j, j))
+            cir = qforte.Circuit(self._Uprep)
 
             cir.add(qforte.gate("H", ancilla_idx, ancilla_idx))
 
@@ -323,11 +321,7 @@ def matrix_element(self, m, n, use_op=False):
                     control_gate_str = "c" + gate_str
                     cV_l.add(qforte.gate(control_gate_str, target, ancilla_idx))
 
-                cir = qforte.Circuit()
-                # TODO (opt): use Uprep
-                for j in range(self._nqb):
-                    if self._ref[j] == 1:
-                        cir.add(qforte.gate("X", j, j))
+                cir = qforte.Circuit(self._Uprep)
 
                 cir.add(qforte.gate("H", ancilla_idx, ancilla_idx))