Tensor from classical action

dev_tools/qualtran_dev_tools/tensor_from_classical_test.py

@@ -0,0 +1,67 @@
+#  Copyright 2024 Google LLC
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#      https://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+from typing import Callable
+
+import numpy as np
+import pytest
+
+from qualtran.simulation.tensor.tensor_from_classical import tensor_from_classical_sim
+from qualtran.symbolics import is_symbolic
+
+from .bloq_finder import get_bloq_examples
+from .tensor_report_card import ExecuteWithTimeout
+
+
+def exec_with_timeout(fn: Callable, *, timeout: float = 10.0):
+    def _run_fn(f, cxn):
+        res, err = None, None
+        try:
+            res = f()
+        except Exception as e:  # pylint: disable=broad-exception-caught
+            err = str(e)
+        cxn.send((res, err))
+
+    runner = ExecuteWithTimeout(timeout=timeout, max_workers=1)
+    runner.submit(_run_fn, {'f': fn})
+    _, output = runner.next_result()
+    return output
+
+
+@pytest.mark.parametrize("be", get_bloq_examples(), ids=lambda be: be.name)
+def test_classical_consistent_with_tensor(be):
+    LIM = 30
+
+    bloq = be.make()
+
+    n = bloq.signature.n_qubits()
+    if is_symbolic(n):
+        pytest.skip(f'symbolic qubits: {n}')
+    if n > LIM:
+        pytest.skip(f'too many qubits: {n=} > {LIM}')
+
+    result = exec_with_timeout(bloq.tensor_contract)
+    if result is None:
+        pytest.skip('timeout: tensor')
+    tensor_direct, err = result
+    if err is not None:
+        pytest.skip(f'no tensor: {err}')
+
+    result = exec_with_timeout(lambda: tensor_from_classical_sim(bloq))
+    if result is None:
+        pytest.skip('timeout: tensor from classical')
+    tensor_classical, err = result
+    if err is not None:
+        pytest.skip(f'no classical action: {err}')
+
+    np.testing.assert_allclose(tensor_classical, tensor_direct, rtol=1e-5, atol=1e-5)

qualtran/bloqs/arithmetic/sorting.py

@@ -84,7 +84,7 @@ def build_composite_bloq(
 
 @bloq_example
 def _comparator() -> Comparator:
-    comparator = Comparator(7)
+    comparator = Comparator(3)
     return comparator
 
 

qualtran/bloqs/basic_gates/swap.py

@@ -299,7 +299,7 @@ def _swap_small() -> Swap:
     return swap_small
 
 
-@bloq_example
+@bloq_example(generalizer=ignore_split_join)
 def _swap_large() -> Swap:
     swap_large = Swap(bitsize=64)
     return swap_large

qualtran/bloqs/basic_gates/swap_test.py

@@ -35,6 +35,7 @@
     _cswap_large,
     _cswap_small,
     _swap,
+    _swap_large,
     _swap_matrix,
     _swap_small,
     Swap,
@@ -224,6 +225,10 @@ def test_swap_small(bloq_autotester):
     bloq_autotester(_swap_small)
 
 
+def test_swap_large(bloq_autotester):
+    bloq_autotester(_swap_large)
+
+
 def test_swap_symb(bloq_autotester):
     if bloq_autotester.check_name == 'serialize':
         pytest.skip("Sympy equality with assumptions.")

qualtran/simulation/tensor/tensor_from_classical.py

@@ -0,0 +1,71 @@
+#  Copyright 2024 Google LLC
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#      https://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+import itertools
+from typing import TYPE_CHECKING
+
+import numpy as np
+from numpy.typing import NDArray
+
+from qualtran import Bloq, Register
+
+if TYPE_CHECKING:
+    from qualtran.simulation.classical_sim import ClassicalValT
+
+
+def _bits_to_classical_reg_data(reg: Register, bits) -> 'ClassicalValT':
+    if reg.shape == ():
+        return reg.dtype.from_bits(bits)
+    return reg.dtype.from_bits_array(np.reshape(bits, reg.shape + (reg.dtype.num_qubits,)))
+
+
+def tensor_from_classical_sim(bloq: Bloq) -> NDArray:
+    left_qubit_counts = tuple(reg.total_bits() for reg in bloq.signature.lefts())
+    left_qubit_splits = np.cumsum(left_qubit_counts)
+
+    n_qubits_left = sum(left_qubit_counts)
+    n_qubits_right = sum(reg.total_bits() for reg in bloq.signature.rights())
+
+    matrix = np.zeros((2,) * (n_qubits_right + n_qubits_left))
+
+    for input_t in itertools.product((0, 1), repeat=n_qubits_left):
+        *inputs_t, last = np.split(input_t, left_qubit_splits)
+        assert np.size(last) == 0
+
+        input_kwargs = {
+            reg.name: _bits_to_classical_reg_data(reg, bits)
+            for reg, bits in zip(bloq.signature.lefts(), inputs_t)
+        }
+        output_args = bloq.call_classically(**input_kwargs)
+
+        if output_args:
+            output_t = np.concatenate(
+                [
+                    reg.dtype.to_bits_array(np.asarray(vals)).flat
+                    for reg, vals in zip(bloq.signature.rights(), output_args)
+                ]
+            )
+        else:
+            output_t = np.array([])
+
+        matrix[tuple([*np.atleast_1d(output_t), *np.atleast_1d(input_t)])] = 1
+
+    shape: tuple[int, ...]
+    if n_qubits_left == 0 and n_qubits_right == 0:
+        shape = ()
+    elif n_qubits_left == 0 or n_qubits_right == 0:
+        shape = (2 ** max(n_qubits_left, n_qubits_right),)
+    else:
+        shape = (2**n_qubits_right, 2**n_qubits_left)
+
+    return matrix.reshape(shape)

qualtran/simulation/tensor/tensor_from_classical_test.py

@@ -0,0 +1,28 @@
+#  Copyright 2024 Google LLC
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#      https://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+import numpy as np
+import pytest
+
+from qualtran import Bloq
+from qualtran.bloqs.basic_gates import TwoBitCSwap, XGate
+
+from .tensor_from_classical import tensor_from_classical_sim
+
+
+@pytest.mark.parametrize("bloq", [XGate(), TwoBitCSwap()], ids=str)
+def test_tensor_consistent_with_classical(bloq: Bloq):
+    from_classical = tensor_from_classical_sim(bloq)
+    from_tensor = bloq.tensor_contract()
+
+    np.testing.assert_allclose(from_classical, from_tensor)