Fix tests, mypy and pylint

Author: tanujkhattar

dev_tools/qualtran_dev_tools/notebook_specs.py

@@ -611,17 +611,17 @@
     # -----   Polynomials defined over Galois Fields (GF) Arithmetic    --------
     # --------------------------------------------------------------------------
     NotebookSpecV2(
-        title='GF Polynomials Split and Join',
+        title='Polynomials over GF($p^m$) - Split and Join',
         module=qualtran.bloqs.gf_poly_arithmetic.gf_poly_split_and_join,
         bloq_specs=[
             qualtran.bloqs.gf_poly_arithmetic.gf_poly_split_and_join._GF_POLY_SPLIT_DOC,
             qualtran.bloqs.gf_poly_arithmetic.gf_poly_split_and_join._GF_POLY_JOIN_DOC,
         ],
     ),
     NotebookSpecV2(
-        title='GF($2^m$) Polynomials Add Constant',
+        title='Polynomials over GF($2^m$) - Add Constant',
         module=qualtran.bloqs.gf_poly_arithmetic.gf2_poly_add_k,
-        bloq_specs=[qualtran.bloqs.gf_poly_arithmetic.gf2_poly_add_k._GF2_ADD_K_DOC],
+        bloq_specs=[qualtran.bloqs.gf_poly_arithmetic.gf2_poly_add_k._GF2_POLY_ADD_K_DOC],
     ),
 ]
 

qualtran/bloqs/gf_poly_arithmetic/__init__.py

@@ -0,0 +1,17 @@
+#  Copyright 2025 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 qualtran.bloqs.gf_poly_arithmetic.gf2_poly_add_k import GF2PolyAddK
+from qualtran.bloqs.gf_poly_arithmetic.gf_poly_split_and_join import GFPolyJoin, GFPolySplit

qualtran/bloqs/gf_poly_arithmetic/gf2_poly_add_k.ipynb

@@ -2,18 +2,18 @@
  "cells": [
   {
    "cell_type": "markdown",
-   "id": "3d480a1d",
+   "id": "71e0dbd8",
    "metadata": {
     "cq.autogen": "title_cell"
    },
    "source": [
-    "# GF($2^m$) Polynomials Add Constant"
+    "# Polynomials over GF($2^m$) - Add Constant"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "218a8310",
+   "id": "c02448cd",
    "metadata": {
     "cq.autogen": "top_imports"
    },
@@ -30,47 +30,48 @@
   },
   {
    "cell_type": "markdown",
-   "id": "203cfaf0",
+   "id": "a5f23a08",
    "metadata": {
-    "cq.autogen": "GF2AddK.bloq_doc.md"
+    "cq.autogen": "GF2PolyAddK.bloq_doc.md"
    },
    "source": [
-    "## `GF2AddK`\n",
-    "In place addition of a constant $k$ for elements in GF($2^m$).\n",
+    "## `GF2PolyAddK`\n",
+    "In place addition of a constant polynomial defined over GF($2^m$).\n",
     "\n",
-    "The bloq implements in place addition of a classical constant $k$ and a quantum register\n",
-    "$|x\\rangle$ storing elements from GF($2^m$). Addition in GF($2^m$) simply reduces to a component\n",
-    "wise XOR, which can be implemented via X gates.\n",
+    "The bloq implements in place addition of a classical constant polynomial $g(x)$ and\n",
+    "a quantum register $|f(x)\\rangle$ storing coefficients of a degree-n polynomial defined\n",
+    "over GF($2^m$). Addition in GF($2^m$) simply reduces to a component wise XOR, which can\n",
+    "be implemented via X gates.\n",
     "\n",
     "$$\n",
-    "|x\\rangle  \\rightarrow |x + k\\rangle\n",
+    "    |f(x)\\rangle  \\rightarrow |f(x) + g(x)\\rangle\n",
     "$$\n",
     "\n",
     "#### Parameters\n",
-    " - `bitsize`: The degree $m$ of the galois field GF($2^m$). Also corresponds to the number of qubits in the input register x.\n",
-    " - `k`: Integer representation of constant over GF($2^m$) that should be added to the input register x. \n",
+    " - `qgf_poly`: An instance of `QGFPoly` type that defines the data type for quantum register $|f(x)\\rangle$ storing coefficients of a degree-n polynomial defined over GF($2^m$).\n",
+    " - `g_x`: An instance of `galois.Poly` that specifies that constant polynomial g(x) defined over GF($2^m$) that should be added to the input register f(x). \n",
     "\n",
     "#### Registers\n",
-    " - `x`: Input THRU register of size $m$ that stores elements from $GF(2^m)$.\n"
+    " - `f_x`: Input THRU register that stores coefficients of polynomial defined over $GF(2^m)$.\n"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "d52372c6",
+   "id": "d40e23be",
    "metadata": {
-    "cq.autogen": "GF2AddK.bloq_doc.py"
+    "cq.autogen": "GF2PolyAddK.bloq_doc.py"
    },
    "outputs": [],
    "source": [
-    "from qualtran.bloqs.gf_arithmetic import GF2AddK"
+    "from qualtran.bloqs.gf_poly_arithmetic import GF2PolyAddK"
    ]
   },
   {
    "cell_type": "markdown",
-   "id": "1b73fd35",
+   "id": "101615d7",
    "metadata": {
-    "cq.autogen": "GF2AddK.example_instances.md"
+    "cq.autogen": "GF2PolyAddK.example_instances.md"
    },
    "source": [
     "### Example Instances"
@@ -79,13 +80,14 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "ab076857",
+   "id": "2cc9165f",
    "metadata": {
-    "cq.autogen": "GF2AddK.gf2_poly_4_8_add_k"
+    "cq.autogen": "GF2PolyAddK.gf2_poly_4_8_add_k"
    },
    "outputs": [],
    "source": [
     "from galois import Poly\n",
+    "\n",
     "from qualtran import QGF, QGFPoly\n",
     "\n",
     "qgf_poly = QGFPoly(4, QGF(2, 3))\n",
@@ -96,26 +98,27 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "991e90de",
+   "id": "9081edee",
    "metadata": {
-    "cq.autogen": "GF2AddK.gf2_poly_add_k_symbolic"
+    "cq.autogen": "GF2PolyAddK.gf2_poly_add_k_symbolic"
    },
    "outputs": [],
    "source": [
     "import sympy\n",
     "from galois import Poly\n",
+    "\n",
     "from qualtran import QGF, QGFPoly\n",
     "\n",
     "n, m = sympy.symbols('n, m', positive=True, integers=True)\n",
     "qgf_poly = QGFPoly(n, QGF(2, m))\n",
-    "gf2_poly_symbolic_add_k = GF2PolyAddK(qgf_poly, Poly([0, 0, 0, 0]))"
+    "gf2_poly_add_k_symbolic = GF2PolyAddK(qgf_poly, Poly([0, 0, 0, 0]))"
    ]
   },
   {
    "cell_type": "markdown",
-   "id": "6a3ecb1e",
+   "id": "11a4d71f",
    "metadata": {
-    "cq.autogen": "GF2AddK.graphical_signature.md"
+    "cq.autogen": "GF2PolyAddK.graphical_signature.md"
    },
    "source": [
     "#### Graphical Signature"
@@ -124,9 +127,9 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "32a100a5",
+   "id": "66b4a211",
    "metadata": {
-    "cq.autogen": "GF2AddK.graphical_signature.py"
+    "cq.autogen": "GF2PolyAddK.graphical_signature.py"
    },
    "outputs": [],
    "source": [
@@ -137,9 +140,9 @@
   },
   {
    "cell_type": "markdown",
-   "id": "3cd8496b",
+   "id": "5940114d",
    "metadata": {
-    "cq.autogen": "GF2AddK.call_graph.md"
+    "cq.autogen": "GF2PolyAddK.call_graph.md"
    },
    "source": [
     "### Call Graph"
@@ -148,9 +151,9 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "605f81e0",
+   "id": "dffd364c",
    "metadata": {
-    "cq.autogen": "GF2AddK.call_graph.py"
+    "cq.autogen": "GF2PolyAddK.call_graph.py"
    },
    "outputs": [],
    "source": [

qualtran/bloqs/gf_poly_arithmetic/gf2_poly_add_k.py

@@ -12,7 +12,7 @@
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
 from functools import cached_property
-from typing import Dict, TYPE_CHECKING
+from typing import Dict, Set, TYPE_CHECKING, Union
 
 import attrs
 import galois
@@ -32,7 +32,7 @@
 
 if TYPE_CHECKING:
     from qualtran import BloqBuilder, Soquet
-    from qualtran.resource_counting import BloqCountDictT, SympySymbolAllocator
+    from qualtran.resource_counting import BloqCountDictT, BloqCountT, SympySymbolAllocator
     from qualtran.simulation.classical_sim import ClassicalValT
 
 
@@ -92,7 +92,9 @@ def build_composite_bloq(self, bb: 'BloqBuilder', *, f_x: 'Soquet') -> Dict[str,
         f_x = bb.add(GFPolyJoin(self.qgf_poly), reg=f_x)
         return {'f_x': f_x}
 
-    def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT':
+    def build_call_graph(
+        self, ssa: 'SympySymbolAllocator'
+    ) -> Union['BloqCountDictT', Set['BloqCountT']]:
         if self.is_symbolic():
             k = ssa.new_symbol('g_x')
             return {GF2AddK(self.qgf_poly.qgf.bitsize, k): self.qgf_poly.degree + 1}
@@ -123,10 +125,10 @@ def _gf2_poly_add_k_symbolic() -> GF2PolyAddK:
 
     n, m = sympy.symbols('n, m', positive=True, integers=True)
     qgf_poly = QGFPoly(n, QGF(2, m))
-    gf2_poly_symbolic_add_k = GF2PolyAddK(qgf_poly, Poly([0, 0, 0, 0]))
-    return gf2_poly_symbolic_add_k
+    gf2_poly_add_k_symbolic = GF2PolyAddK(qgf_poly, Poly([0, 0, 0, 0]))
+    return gf2_poly_add_k_symbolic
 
 
-_GF2_ADD_K_DOC = BloqDocSpec(
-    bloq_cls=GF2AddK, examples=(_gf2_poly_4_8_add_k, _gf2_poly_add_k_symbolic)
+_GF2_POLY_ADD_K_DOC = BloqDocSpec(
+    bloq_cls=GF2PolyAddK, examples=(_gf2_poly_4_8_add_k, _gf2_poly_add_k_symbolic)
 )

qualtran/bloqs/gf_poly_arithmetic/gf2_poly_add_k_test.py

@@ -11,6 +11,7 @@
 #  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
 from galois import Poly
 
 from qualtran.bloqs.gf_poly_arithmetic.gf2_poly_add_k import (
@@ -31,12 +32,14 @@ def test_gf2_poly_symbolic_add_k(bloq_autotester):
 def test_gf2_poly_add_k_classical_sim():
     bloq = _gf2_poly_4_8_add_k.make()
     f_x = Poly(bloq.qgf_poly.qgf.gf_type([0, 1, 2, 3, 4]))
-    assert bloq.call_classically(f_x=f_x)[0] == f_x + bloq.g_x
-
-    f_x_range = [
-        Poly(bloq.qgf_poly.qgf.gf_type([0, 0, 0, 0, 0])),
-        Poly(bloq.qgf_poly.qgf.gf_type([7, 7, 7, 7, 7])),
-        Poly(bloq.qgf_poly.qgf.gf_type([0, 0, 3, 5, 7])),
-        Poly(bloq.qgf_poly.qgf.gf_type([2, 3, 5, 0, 0])),
-    ]
+    assert bloq.call_classically(f_x=f_x)[0] == f_x + bloq.g_x  # type: ignore[arg-type]
+
+    f_x_range = np.asarray(
+        [
+            Poly(bloq.qgf_poly.qgf.gf_type([0, 0, 0, 0, 0])),
+            Poly(bloq.qgf_poly.qgf.gf_type([7, 7, 7, 7, 7])),
+            Poly(bloq.qgf_poly.qgf.gf_type([0, 0, 3, 5, 7])),
+            Poly(bloq.qgf_poly.qgf.gf_type([2, 3, 5, 0, 0])),
+        ]
+    )
     assert_consistent_classical_action(bloq, f_x=f_x_range)