fix _mul_via_repeated_add (Fxp shifting is buggy)

qualtran/bloqs/rotations/phase_gradient.py

@@ -366,37 +366,30 @@ def _add_into_phase_grad() -> AddIntoPhaseGrad:
 
 
 def _fxp(x: float, n: 'SymbolicInt') -> Fxp:
-    """When 0 <= x < 1, constructs an n-bit fixed point representation with nice properties.
-
-    Specifically,
-    -   op_sizing='same' and const_op_sizing='same' ensure that the returned object is not resized
-        to a bigger fixed point number when doing operations with other Fxp objects.
-    -   shifting='trunc' ensures that when shifting the Fxp integer to left / right; the digits are
-        truncated and no rounding occurs
-    -   overflow='wrap' ensures that when performing operations where result overflows, the overflowed
-        digits are simply discarded.
-    """
+    """When 0 <= x < 1, constructs an n-bit fixed point representation with nice properties."""
     assert 0 <= x < 1
-    return Fxp(
-        x,
-        dtype=f'fxp-u{n}/{n}',
-        op_sizing='same',
-        const_op_sizing='same',
-        shifting='trunc',
-        overflow='wrap',
-    )
+    return QFxp(n, n).float_to_fxp(x, require_exact=False)
 
 
 def _mul_via_repeated_add(x_fxp: Fxp, gamma_fxp: Fxp, out: int) -> Fxp:
     """Multiplication via repeated additions algorithm described in Appendix D5"""
 
     res = _fxp(0, out)
+    x_fxp = x_fxp.copy()
+    x_fxp.resize(n_int=x_fxp.n_int, n_frac=max(out, x_fxp.n_frac))
     for i, bit in enumerate(gamma_fxp.bin()):
         if bit == '0':
             continue
         shift = gamma_fxp.n_int - i - 1
         # Left/Right shift by `shift` bits.
-        res += x_fxp << shift if shift > 0 else x_fxp >> abs(shift)
+        # Issue: Fxp << and >> does not preserve config, so we use `* 2**shift` instead.
+        x_shifted = x_fxp.copy()
+        if shift > 0:
+            x_shifted.set_val(x_fxp.val << np.array(shift, dtype=x_fxp.val.dtype), raw=True)
+        else:
+            x_shifted.set_val(x_fxp.val >> np.array(-shift, dtype=x_fxp.val.dtype), raw=True)
+        assert x_shifted.overflow == 'wrap' and x_shifted.shifting == 'trunc'
+        res += x_shifted
     return res
 
 

qualtran/bloqs/rotations/phasing_via_cost_function_test.py

@@ -127,7 +127,6 @@ def test_hamming_weight_phasing_using_phase_via_cost_function(
 @attrs.frozen
 class TestSquarePhasing(GateWithRegisters):
     bitsize: int
-    normalize: bool
     gamma: float
     eps: float
     use_phase_gradient: bool = False
@@ -138,9 +137,7 @@ def signature(self) -> 'Signature':
 
     @property
     def cost_reg(self) -> Register:
-        return Register(
-            'result', QFxp(2 * self.bitsize, 2 * self.bitsize * int(self.normalize), signed=False)
-        )
+        return Register('result', QFxp(2 * self.bitsize, 2 * self.bitsize, signed=False))
 
     @property
     def phase_gradient_oracle(self) -> QvrPhaseGradient:
@@ -169,30 +166,17 @@ def build_composite_bloq(self, bb: 'BloqBuilder', **soqs: 'SoquetT') -> Dict[str
         return soqs
 
 
-@pytest.mark.slow
-@pytest.mark.parametrize('normalize_cost_function', [True, False])
-@pytest.mark.parametrize('use_phase_gradient', [True, False])
-@pytest.mark.parametrize('gamma, eps', [(0.1, 5e-2), (1.20345, 5e-2), (-1.1934341, 5e-2)])
+@pytest.mark.parametrize('use_phase_gradient', [pytest.param(True, marks=pytest.mark.slow), False])
+@pytest.mark.parametrize('gamma, eps', [(0.125, 5e-2), (1.1875, 5e-2), (-1.375, 5e-2)])
 @pytest.mark.parametrize('n', [2])
 def test_square_phasing_via_phase_gradient(
-    n: int, gamma: float, eps: float, use_phase_gradient: bool, normalize_cost_function: bool
+    n: int, gamma: float, eps: float, use_phase_gradient: bool
 ):
     initial_state = np.array([1 / np.sqrt(2**n)] * 2**n)
-    normalization_factor = 1 if normalize_cost_function else 4**n
-    phases = np.array(
-        [
-            np.exp(1j * 2 * np.pi * gamma * x**2 * normalization_factor / 4**n)
-            for x in range(2**n)
-        ]
-    )
+    phases = np.array([np.exp(1j * 2 * np.pi * gamma * x**2 / 4**n) for x in range(2**n)])
     expected_final_state = np.multiply(initial_state, phases)
-    test_bloq_one = TestSquarePhasing(
-        n, True, gamma * normalization_factor, eps, use_phase_gradient
-    )
-    test_bloq_two = TestSquarePhasing(
-        n, False, gamma * normalization_factor / (4**n), eps, use_phase_gradient
-    )
-    for test_bloq in [test_bloq_one, test_bloq_two]:
+    test_bloq = TestSquarePhasing(n, gamma, eps, use_phase_gradient)
+    for test_bloq in [test_bloq]:
         bb = BloqBuilder()
         a = bb.allocate(n)
         a = bb.add(OnEach(n, Hadamard()), q=a)

qualtran/bloqs/rotations/quantum_variable_rotation.py

@@ -240,7 +240,8 @@ def find_optimal_phase_grad_size(gamma_fxp: Fxp, cost_dtype: QFxp, eps: float) -
 
     def is_good_phase_grad_size(phase_bitsize: int):
         res = _mul_via_repeated_add(cost_fxp, gamma_fxp, phase_bitsize)
-        return np.abs(res.get_val() - expected_val) <= eps
+        actual_val = res.get_val() % 1
+        return np.abs(actual_val - expected_val) <= eps
 
     low, high, ans = 1, 100, None
     while low <= high:
@@ -464,7 +465,7 @@ def gamma_fxp(self) -> Fxp:
         if is_symbolic(self.gamma):
             raise ValueError(f"Cannot compute Fxp from symbolic {self.gamma=}")
 
-        return self.gamma_dtype.float_to_fxp(abs(self.gamma))
+        return self.gamma_dtype.float_to_fxp(abs(self.gamma), require_exact=False)
 
     @cached_property
     def gamma_dtype(self) -> QFxp: