address comments

qualtran/bloqs/mod_arithmetic/mod_division.py

@@ -449,6 +449,17 @@ def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT':
     def on_classical_vals(
         self, u: int, v: int, r: int, s: int, m: int, f: int
     ) -> Dict[str, 'ClassicalValT']:
+        """This is a classical encoding of figure 15 of https://arxiv.org/pdf/2302.06639.
+
+        The variables `m` and the local variables `a` and `b` translate into evaluating the if
+        conditions in `Algorithm 2 `. The meaning of the variables are:
+            - `a`: is `u` even?
+            - `b`: are both `u` and `v` even?
+            - `m`: is `u` odd and `v` even?
+            - `f`: classically once `f = 0` the algorithm terminates.
+        `a` and `b` are local and cleaned after each iteration. The variable `m` is kept and
+        is used in uncomputation.
+        """
         a = b = 0
         assert m == 0
         m ^= f & (v == 0)

qualtran/bloqs/mod_arithmetic/mod_division_test.py

@@ -18,6 +18,7 @@
 import sympy
 
 import qualtran.testing as qlt_testing
+from qualtran import QMontgomeryUInt
 from qualtran.bloqs.mod_arithmetic.mod_division import _kaliskimodinverse_example, KaliskiModInverse
 from qualtran.resource_counting import get_cost_value, QECGatesCost
 from qualtran.resource_counting.generalizers import ignore_alloc_free, ignore_split_join
@@ -28,17 +29,17 @@
 def test_kaliski_mod_inverse_classical_action(bitsize, mod):
     blq = KaliskiModInverse(bitsize, mod)
     cblq = blq.decompose_bloq()
-    p2 = pow(2, bitsize, mod)
+    dtype = QMontgomeryUInt(bitsize)
+    R = pow(2, bitsize, mod)
     for x in range(1, mod):
         if math.gcd(x, mod) != 1:
             continue
-        x_montgomery = (x * p2) % mod
-        inv_x = pow(x, -1, mod)
-        inv_x_montgomery = (inv_x * p2) % mod
+        x_montgomery = dtype.uint_to_montgomery(x, mod)
         res = blq.call_classically(x=x_montgomery)
         assert res == cblq.call_classically(x=x_montgomery)
         assert len(res) == 2
-        assert res[0] == inv_x_montgomery
+        assert res[0] == dtype.montgomery_inverse(x_montgomery, mod)
+        assert dtype.montgomery_product(res[0], x_montgomery, mod) == R
 
 
 @pytest.mark.parametrize('bitsize', [5, 6])