Repository: quantumlib/Qualtran

Issue:
Incorrect `on_classical_vals` overrides for non classical bloqs
I ran into this while implementing #1514. The tensors of these bloq examples seem to have non-trivial phases, which makes them non-classical. The following bloq examples fail in #1514 (I've restricted to bloqs with up to 9 qubits):
```
sub_from_small

E Mismatched elements: 144 / 65536 (0.22%)
E Max absolute difference: 1.
E Max relative difference: 1.
E x: array([[1., 0., 0., ..., 0., 0., 0.],
E [0., 1., 0., ..., 0., 0., 0.],
E [0., 0., 1., ..., 0., 0., 0.],...
E y: array([[1.+0.j, 0.+0.j, 0.+0.j, ..., 0.+0.j, 0.+0.j, 0.+0.j],
E [0.+0.j, 1.+0.j, 0.+0.j, ..., 0.+0.j, 0.+0.j, 0.+0.j],
E [0.+0.j, 0.+0.j, 1.+0.j, ..., 0.+0.j, 0.+0.j, 0.+0.j],...
```
```
cmod_add_k_small

E Mismatched elements: 18 / 1024 (1.76%)
E Max absolute difference: 1.
E Max relative difference: 0.
E x: array([[1., 0., 0., ..., 0., 0., 0.],
E [0., 1., 0., ..., 0., 0., 0.],
E [0., 0., 1., ..., 0., 0., 0.],...
E y: array([[1.+0.j, 0.+0.j, 0.+0.j, ..., 0.+0.j, 0.+0.j, 0.+0.j],
E [0.+0.j, 1.+0.j, 0.+0.j, ..., 0.+0.j, 0.+0.j, 0.+0.j],
E [0.+0.j, 0.+0.j, 1.+0.j, ..., 0.+0.j, 0.+0.j, 0.+0.j],...
```
```
moddbl_small

E Mismatched elements: 3 / 256 (1.17%)
E Max absolute difference: 1.
E Max relative difference: 1.
E x: array([[1., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
E [0., 0., 0., 0., 0., 0., 0., 1., 0., 0., 0., 0., 0., 0., 0., 0.],
E [0., 1., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],...
E y: array([[1.+0.j, 0.+0.j, 0.+0.j, 0.+0.j, 0.+0.j, 0.+0.j, 0.+0.j, 0.+0.j,
E 0.+0.j, 0.+0.j, 0.+0.j, 0.+0.j, 0.+0.j, 0.+0.j, 0.+0.j, 0.+0.j],
E [0.+0.j, 0.+0.j, 0.+0.j, 0.+0.j, 0.+0.j, 0.+0.j, 0.+0.j, 1.+0.j,...
```
```
approx_cswap_small

E Mismatched elements: 120 / 262144 (0.0458%)
E Max absolute difference: 2.
E Max relative difference: 2.
E x: array([[1., 0., 0., ..., 0., 0., 0.],
E [0., 1., 0., ..., 0., 0., 0.],
E [0., 0., 1., ..., 0., 0., 0.],...
E y: array([[1.000000e+00+6.471125e-32j, 0.000000e+00+0.000000e+00j,
E 0.000000e+00+0.000000e+00j, ..., 0.000000e+00+0.000000e+00j,
E 0.000000e+00+0.000000e+00j, 0.000000e+00+0.000000e+00j],...
```

Discussion:
@mpharrigan: Do you know where the non-trivial phases are being introduced? Presumably if you write a decomposition of e.g. ModDbl using only classical operations (whose matrices only have `+1` entries) then you can't get any phases popping up, right? @tanujkhattar: fyi -- `approx_cswap_small`'s tensor would have a non-trivial phases since it implements an approximate cswap which is correct upto a phase, in order to have 4 T count instead of 7,  which gets inverted when the approx cswap inverse is applied. This is similar to the 4 T version of Toffoli which applies a toffoli upto a relative phase (we don't explicitly have a bloq for this yet)@mpharrigan: Yeah, the approx cswap is clearer. It's got "approx" in the name, so it's less surprising that there are phases. The other ones are surprising to me@NoureldinYosri: `sub_from_small ` => the classical action of `SubtractFrom` is wrong. it uses `fmod` which gives wrong answer on overflow instead it should call `simulation.classical_sim.add_ints` ... this also means that the classical action tests for this bloq are insufficient.

cmod_add_k_small and moddbl_small are because their constructions assume correct range ... that they assume they will be called on states $x < \mod < 2^\mathrm{bitsize}$ ... outside this range (i.e. $\mod \leq x < 2^\mathrm{bitsize}$) which happens for tensor contraction their constructions doen't work .... I can change the `on_classical_vals` to take this into account (if input is outside the valid range return it as is) ... this will make the tensor contraction work but that may introduce a difference in behaviour between the bloq and it decomposition for those out of bound inputs@mpharrigan: is it obvious how those corner cases lead to _phases_ in the tensors? @anurudhp: I looked at the failing tests again, and I realized that the arithmetic bloqs fail because the inputs outside range (i.e. `>= mod`) do not match (the max abs diff is `1`, which is from `1`s and `0`s mismatching).

Only the `cswap_approx` fails because of a phase mismatch.@NoureldinYosri: > is it obvious how those corner cases lead to _phases_ in the tensors?

there are no phases except in `cswap_approx`. as anurudhp this is mismatch is between 0 and 1 for inputs outside the valid range for the bloqs@mpharrigan: ok, that makes way more sense 

Solution patch:
diff --git a/qualtran/bloqs/arithmetic/subtraction.py b/qualtran/bloqs/arithmetic/subtraction.py
index ef5cb519..f0625570 100644
--- a/qualtran/bloqs/arithmetic/subtraction.py
+++ b/qualtran/bloqs/arithmetic/subtraction.py
@@ -12,7 +12,6 @@
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
 
-import math
 from typing import Dict, Optional, Tuple, TYPE_CHECKING, Union
 
 import numpy as np
@@ -40,6 +39,7 @@ from qualtran.bloqs.basic_gates import OnEach, XGate
 from qualtran.bloqs.bookkeeping import Allocate, Cast, Free
 from qualtran.bloqs.mcmt.multi_target_cnot import MultiTargetCNOT
 from qualtran.drawing import Text
+from qualtran.simulation.classical_sim import add_ints
 
 if TYPE_CHECKING:
     from qualtran.drawing import WireSymbol
@@ -270,10 +270,15 @@ class SubtractFrom(Bloq):
     def on_classical_vals(
         self, a: 'ClassicalValT', b: 'ClassicalValT'
     ) -> Dict[str, 'ClassicalValT']:
-        unsigned = isinstance(self.dtype, (QUInt, QMontgomeryUInt))
-        bitsize = self.dtype.bitsize
-        N = 2**bitsize if unsigned else 2 ** (bitsize - 1)
-        return {'a': a, 'b': int(math.fmod(b - a, N))}
+        return {
+            'a': a,
+            'b': add_ints(
+                int(b),
+                -int(a),
+                num_bits=int(self.dtype.bitsize),
+                is_signed=isinstance(self.dtype, QInt),
+            ),
+        }
 
     def wire_symbol(
         self, reg: Optional['Register'], idx: Tuple[int, ...] = tuple()
diff --git a/qualtran/bloqs/mod_arithmetic/mod_addition.py b/qualtran/bloqs/mod_arithmetic/mod_addition.py
index cb7ab639..db4fe7db 100644
--- a/qualtran/bloqs/mod_arithmetic/mod_addition.py
+++ b/qualtran/bloqs/mod_arithmetic/mod_addition.py
@@ -87,7 +87,17 @@ class ModAdd(Bloq):
     def on_classical_vals(
         self, x: 'ClassicalValT', y: 'ClassicalValT'
     ) -> Dict[str, 'ClassicalValT']:
-        return {'x': x, 'y': (x + y) % self.mod}
+        if not (0 <= x < self.mod):
+            raise ValueError(
+                f'{x=} is outside the valid interval for modular addition [0, {self.mod})'
+            )
+        if not (0 <= y < self.mod):
+            raise ValueError(
+                f'{y=} is outside the valid interval for modular addition [0, {self.mod})'
+            )
+
+        y = (x + y) % self.mod
+        return {'x': x, 'y': y}
 
     def build_composite_bloq(self, bb: 'BloqBuilder', x: Soquet, y: Soquet) -> Dict[str, 'SoquetT']:
         if is_symbolic(self.bitsize):
@@ -307,6 +317,12 @@ class CModAddK(Bloq):
             return {'ctrl': 0, 'x': x}
 
         assert ctrl == 1, 'Bad ctrl value.'
+
+        if not (0 <= x < self.mod):
+            raise ValueError(
+                f'{x=} is outside the valid interval for modular addition [0, {self.mod})'
+            )
+
         x = (x + self.k) % self.mod
         return {'ctrl': ctrl, 'x': x}
 
@@ -492,7 +508,17 @@ class CModAdd(Bloq):
         if ctrl != self.cv:
             return {'ctrl': ctrl, 'x': x, 'y': y}
 
-        return {'ctrl': ctrl, 'x': x, 'y': (x + y) % self.mod}
+        if not (0 <= x < self.mod):
+            raise ValueError(
+                f'{x=} is outside the valid interval for modular addition [0, {self.mod})'
+            )
+        if not (0 <= y < self.mod):
+            raise ValueError(
+                f'{y=} is outside the valid interval for modular addition [0, {self.mod})'
+            )
+
+        y = (x + y) % self.mod
+        return {'ctrl': ctrl, 'x': x, 'y': y}
 
     def build_composite_bloq(
         self, bb: 'BloqBuilder', ctrl, x: Soquet, y: Soquet


Test patch:
diff --git a/qualtran/bloqs/arithmetic/subtraction_test.py b/qualtran/bloqs/arithmetic/subtraction_test.py
index 40e2bc6f..fe9a5cb4 100644
--- a/qualtran/bloqs/arithmetic/subtraction_test.py
+++ b/qualtran/bloqs/arithmetic/subtraction_test.py
@@ -160,3 +160,12 @@ def test_subtract_from_bloq_decomposition():
         want[(a << 4) | c][a_b] = 1
     got = gate.tensor_contract()
     np.testing.assert_allclose(got, want)
+
+
+@pytest.mark.parametrize('bitsize', range(2, 5))
+def test_subtractfrom_classical_action(bitsize):
+    dtype = QInt(bitsize)
+    blq = SubtractFrom(dtype)
+    qlt_testing.assert_consistent_classical_action(
+        blq, a=tuple(dtype.get_classical_domain()), b=tuple(dtype.get_classical_domain())
+    )
diff --git a/qualtran/bloqs/mod_arithmetic/mod_addition_test.py b/qualtran/bloqs/mod_arithmetic/mod_addition_test.py
index 455670bb..bd5a11f4 100644
--- a/qualtran/bloqs/mod_arithmetic/mod_addition_test.py
+++ b/qualtran/bloqs/mod_arithmetic/mod_addition_test.py
@@ -208,3 +208,15 @@ def test_cmod_add_complexity_vs_ref():
 def test_mod_add_classical_action(bitsize, prime):
     b = ModAdd(bitsize, prime)
     assert_consistent_classical_action(b, x=range(prime), y=range(prime))
+
+
+def test_cmodadd_tensor():
+    blq = CModAddK(bitsize=4, mod=7, k=1)
+    want = np.zeros((7, 7))
+    for i in range(7):
+        j = (i + 1) % 7
+        want[j, i] = 1
+
+    tn = blq.tensor_contract()
+    np.testing.assert_allclose(tn[:7, :7], np.eye(7))  # ctrl = 0
+    np.testing.assert_allclose(tn[16 : 16 + 7, 16 : 16 + 7], want)  # ctrl = 1