Suppress new Pyrefly findings before upgrading to v1.2.0

.github/workflows/build_and_test.yml

@@ -32,14 +32,15 @@ jobs:
 
       - name: "Run `bazel build`"
         run: |
-          bazel build -c opt //...
+          bazel --host_jvm_args=-Xmx32g build --jobs=8 -c opt //...
 
       - name: "Run `bazel test`"
         run: |
-          bazel test -c opt \
+          bazel --host_jvm_args=-Xmx32g test -c opt \
             --local_cpu_resources=4 \
             --test_output=errors \
             --test_size_filters=small \
             --test_timeout=1800 \
             --experimental_ui_max_stdouterr_bytes=10485760 \
+            --jobs=8 \
             //...

.github/workflows/periodic_tests.yml

@@ -19,14 +19,15 @@ jobs:
 
       - name: "Run `bazel build`"
         run: |
-          bazel build //...
+          bazel --host_jvm_args=-Xmx32g build --jobs=8 //...
 
       - name: "Run `bazel test`"
         run: |
-          bazel test \
+          bazel test --host_jvm_args=-Xmx32g \
             --local_cpu_resources=4 \
             --test_output=errors \
             --test_size_filters=medium,large \
             --test_timeout=3600 \
             --experimental_ui_max_stdouterr_bytes=10485760 \
+            --jobs=8 \
             //...

BUILD

@@ -96,6 +96,10 @@ py_library(
             "jaxite_ckks/*",
         ],
     ),
+    data = [
+        "jaxite_ec/configurations.toml",
+        # "@jaxite//jaxite_ec/c_kernels:distribution.so",
+    ],
     visibility = [":internal"],
     deps = [
         # copybara: xprof_analysis_client  # buildcleaner: keep
@@ -106,7 +110,10 @@ py_library(
         "@jaxite_deps//jaxlib",
         # copybara: jax/experimental:pallas_lib
         # copybara: jax/experimental:pallas_tpu
+        "@jaxite//jaxite_ec/c_kernels:build",
         "@jaxite_deps//numpy",
+        # copybara: pandas
+        # copybara: toml
     ],
 )
 
@@ -160,11 +167,21 @@ tpu_test(
 )
 
 tpu_test(
-    name = "jaxite_ec_finite_field_test",
-    size = "large",
-    timeout = "moderate",
+    name = "ec_finite_field_test",
     srcs = ["jaxite_ec/finite_field_test.py"],
-    shard_count = 3,
+    deps = [
+        ":jaxite",
+        "@abseil-py//absl/testing:absltest",
+        "@abseil-py//absl/testing:parameterized",
+        "@jaxite_deps//jax",
+        "@jaxite_deps//jaxlib",
+        "@jaxite_deps//numpy",
+    ],
+)
+
+tpu_test(
+    name = "ec_finite_field_perf_test",
+    srcs = ["jaxite_ec/finite_field_perf_test.py"],
     deps = [
         ":jaxite",
         # copybara: xprof_analysis_client  # buildcleaner: keep
@@ -174,54 +191,60 @@ tpu_test(
         "@jaxite_deps//jax",
         "@jaxite_deps//jaxlib",
         "@jaxite_deps//numpy",
+        # copybara: toml
     ],
 )
 
 tpu_test(
-    name = "msm_test",
-    size = "large",
-    timeout = "eternal",
-    srcs = [
-        "jaxite_ec/msm_test.py",
+    name = "elliptic_curve_test",
+    srcs = ["jaxite_ec/elliptic_curve_test.py"],
+    deps = [
+        ":jaxite",
+        "@abseil-py//absl/testing:absltest",
+        "@abseil-py//absl/testing:parameterized",
+        "@jaxite_deps//jax",
+        "@jaxite_deps//jaxlib",
+        "@jaxite_deps//numpy",
+        # copybara: toml
     ],
-    data = [
-        "jaxite_ec/test_case/t1/zprize_msm_curve_377_bases_dim_1_seed_0.csv",
-        "jaxite_ec/test_case/t1/zprize_msm_curve_377_res_dim_1_seed_0.csv",
-        "jaxite_ec/test_case/t1/zprize_msm_curve_377_scalars_dim_1_seed_0.csv",
-        "jaxite_ec/test_case/t1024/zprize_msm_curve_377_bases_dim_1024_seed_0.csv",
-        "jaxite_ec/test_case/t1024/zprize_msm_curve_377_res_dim_1024_seed_0.csv",
-        "jaxite_ec/test_case/t1024/zprize_msm_curve_377_scalars_dim_1024_seed_0.csv",
-        "jaxite_ec/test_case/t2/zprize_msm_curve_377_bases_dim_2_seed_0.csv",
-        "jaxite_ec/test_case/t2/zprize_msm_curve_377_res_dim_2_seed_0.csv",
-        "jaxite_ec/test_case/t2/zprize_msm_curve_377_scalars_dim_2_seed_0.csv",
-        "jaxite_ec/test_case/t4/zprize_msm_curve_377_bases_dim_4_seed_0.csv",
-        "jaxite_ec/test_case/t4/zprize_msm_curve_377_res_dim_4_seed_0.csv",
-        "jaxite_ec/test_case/t4/zprize_msm_curve_377_scalars_dim_4_seed_0.csv",
-        "jaxite_ec/test_case/t8/zprize_msm_curve_377_bases_dim_8_seed_0.csv",
-        "jaxite_ec/test_case/t8/zprize_msm_curve_377_res_dim_8_seed_0.csv",
-        "jaxite_ec/test_case/t8/zprize_msm_curve_377_scalars_dim_8_seed_0.csv",
+)
+
+tpu_test(
+    name = "elliptic_curve_perf_test",
+    srcs = ["jaxite_ec/elliptic_curve_perf_test.py"],
+    deps = [
+        ":jaxite",
+        # copybara: xprof_analysis_client  # buildcleaner: keep
+        # copybara: xprof_session  # buildcleaner: keep
+        "@abseil-py//absl/testing:absltest",
+        "@abseil-py//absl/testing:parameterized",
+        "@jaxite_deps//jax",
+        "@jaxite_deps//jaxlib",
+        "@jaxite_deps//numpy",
+        # copybara: toml
     ],
-    shard_count = 3,
-    tags = ["manual"],
+)
+
+tpu_test(
+    name = "multiscalar_multiplication_test",
+    srcs = ["jaxite_ec/multiscalar_multiplication_test.py"],
+    data = glob(["jaxite_ec/data/t1024/*.csv"]),
     deps = [
         ":jaxite",
         # copybara: xprof_analysis_client  # buildcleaner: keep
         # copybara: xprof_session  # buildcleaner: keep
-        # copybara: resources
         "@abseil-py//absl/testing:absltest",
         "@abseil-py//absl/testing:parameterized",
         "@jaxite_deps//jax",
         "@jaxite_deps//jaxlib",
         "@jaxite_deps//numpy",
+        # copybara: toml
     ],
 )
 
 tpu_test(
-    name = "elliptic_curve_test",
-    size = "large",
-    timeout = "long",
-    srcs = ["jaxite_ec/elliptic_curve_test.py"],
-    shard_count = 16,
+    name = "multiscalar_multiplication_perf_test",
+    srcs = ["jaxite_ec/multiscalar_multiplication_perf_test.py"],
     deps = [
         ":jaxite",
         # copybara: xprof_analysis_client  # buildcleaner: keep
@@ -234,6 +257,34 @@ tpu_test(
     ],
 )
 
+tpu_test(
+    name = "number_theory_transform_test",
+    srcs = ["jaxite_ec/number_theory_transform_test.py"],
+    deps = [
+        ":jaxite",
+        "@abseil-py//absl/testing:absltest",
+        "@abseil-py//absl/testing:parameterized",
+        "@jaxite_deps//jax",
+        "@jaxite_deps//jaxlib",
+        "@jaxite_deps//numpy",
+    ],
+)
+
+tpu_test(
+    name = "number_theory_transform_perf_test",
+    size = "large",
+    timeout = "eternal",
+    srcs = ["jaxite_ec/number_theory_transform_perf_test.py"],
+    deps = [
+        ":jaxite",
+        "@abseil-py//absl/testing:absltest",
+        "@abseil-py//absl/testing:parameterized",
+        "@jaxite_deps//jax",
+        "@jaxite_deps//jaxlib",
+        "@jaxite_deps//numpy",
+    ],
+)
+
 cpu_gpu_tpu_test(
     name = "decomposition_test",
     size = "small",

MODULE.bazel

@@ -4,20 +4,21 @@ bazel_dep(name = "bazel_skylib", version = "1.9.0")
 bazel_dep(name = "rules_license", version = "1.0.0")
 bazel_dep(name = "abseil-py", version = "2.1.0")
 bazel_dep(name = "rules_python", version = "1.5.1")
+bazel_dep(name = "rules_cc", version = "0.2.18")
 
 # Hermetic python setup
 python = use_extension("@rules_python//python/extensions:python.bzl", "python")
 
 pip = use_extension("@rules_python//python/extensions:pip.bzl", "pip")
 
-python.toolchain(python_version = "3.11")
+python.toolchain(python_version = "3.13")
 
 pip.parse(
     hub_name = "jaxite_deps",
-    python_version = "3.11",
+    python_version = "3.13",
     requirements_lock = "//:requirements.txt",
 )
 
-use_repo(python, python_3_11 = "python_3_11")
+use_repo(python, python_3_13 = "python_3_13")
 
 use_repo(pip, "jaxite_deps")

jaxite/jaxite_cggi/bootstrap.py

@@ -88,11 +88,11 @@ def gen_bootstrapping_key(
           k,
           rlwe_sk.modulus_degree,
       ),
-      dtype=jnp.uint32,
+      dtype=jnp.uint32,  # pyrefly: ignore[bad-argument-type]
   )
   error_samples = prg.rounded_normal(
       shape=(padded_num_bsk_encryptions, num_blocks + 1, levels),
-      dtype=jnp.uint32,
+      dtype=jnp.uint32,  # pyrefly: ignore[bad-argument-type]
   )
 
   lwe_sk_data = lwe_sk.key_data.astype(jnp.uint32)
@@ -614,7 +614,7 @@ def sample_extract(ciphertext: rlwe.RlweCiphertext) -> types.LweCiphertext:
 @jax.named_call
 @functools.partial(jax.jit, static_argnames="poly_deg")
 def jit_sample_extract(
-    rlwe_ciphertext: jnp.ndarray, poly_deg: jnp.uint32
+    rlwe_ciphertext: jnp.ndarray, poly_deg: jnp.uint32  # pyrefly: ignore[not-a-type]
 ) -> types.LweCiphertext:
   """Extracts an LWE encryption of the constant term encrypted by the input.
 

jaxite/jaxite_cggi/decomposition.py

@@ -32,7 +32,7 @@ class DecompositionParameters:
     jax.jit, static_argnames=("base_log", "num_levels", "total_bit_length")
 )
 def decompose(
-    x: jnp.uint32,
+    x: jnp.uint32,  # pyrefly: ignore[not-a-type]
     base_log: int = 4,
     num_levels: int = 3,
     total_bit_length: int = 32,
@@ -92,7 +92,7 @@ def decompose(
     jax.jit, static_argnames=("base_log", "num_levels", "total_bit_length")
 )
 def recomposition_summands(
-    x: jnp.uint32,
+    x: jnp.uint32,  # pyrefly: ignore[not-a-type]
     base_log: int = 4,
     num_levels: int = 3,
     total_bit_length: int = 32,
@@ -130,7 +130,7 @@ def recomposition_summands(
 @functools.partial(jax.jit, static_argnames=("base_log", "total_bit_length"))
 def recompose(
     digits: DecomposedInt, base_log: int = 4, total_bit_length: int = 32
-) -> jnp.uint32:
+) -> jnp.uint32:  # pyrefly: ignore[not-a-type]
   """The inverse of decompose.
 
   Note num_levels == len(digits).
@@ -153,8 +153,8 @@ def recompose(
     jax.jit, static_argnames=("base_log", "num_levels", "total_bit_length")
 )
 def signed_decomposition(
-    x: jnp.uint32,
-    base_log: jnp.uint32,
+    x: jnp.uint32,  # pyrefly: ignore[not-a-type]
+    base_log: jnp.uint32,  # pyrefly: ignore[not-a-type]
     num_levels: int,
     total_bit_length: int = 32,
 ) -> jnp.ndarray:

jaxite/jaxite_cggi/encoding.py

@@ -67,7 +67,7 @@ def encode(
     message: Union[types.LweCleartext, jnp.ndarray],
     params: EncodingParameters,
     test_polynomial_encoding: bool = False,
-) -> Union[types.LwePlaintext, jnp.ndarray]:
+) -> Union[types.LwePlaintext, jnp.ndarray]:  # pyrefly: ignore[not-a-type]
   """Encode a plaintext or array of plaintexts for use in a TFHE ciphertext.
 
   The bits are organized so that the top bits are padding for overflow,
@@ -109,7 +109,7 @@ def encode(
 
 
 def decode_without_removing_padding(
-    plaintext: types.LwePlaintext, params: EncodingParameters
+    plaintext: types.LwePlaintext, params: EncodingParameters  # pyrefly: ignore[not-a-type]
 ) -> types.LweCleartext:
   """Decode a plaintext without removing padding.
 
@@ -128,7 +128,7 @@ def decode_without_removing_padding(
 
 
 def decode(
-    plaintext: types.LwePlaintext, params: EncodingParameters
+    plaintext: types.LwePlaintext, params: EncodingParameters  # pyrefly: ignore[not-a-type]
 ) -> types.LweCleartext:
   """Decode a plaintext.
 
@@ -176,7 +176,7 @@ def round_to_power_of_2(arr: jnp.ndarray, log_pow_of_2: int) -> jnp.ndarray:
 
 
 def extract_noise(
-    plaintext: types.LwePlaintext, encoding_params: EncodingParameters
+    plaintext: types.LwePlaintext, encoding_params: EncodingParameters  # pyrefly: ignore[not-a-type]
 ) -> int:
   """Extracts the noise bits of a plaintext as a (signed) int."""
   rounded = remove_noise(plaintext, encoding_params)

jaxite/jaxite_cggi/encoding_test.py

@@ -52,7 +52,7 @@ def test_encode_decode_message_bounds_succeeds(
         total_bit_length=32, message_bit_length=16, padding_bit_length=0
     )
 
-    encoded: types.LwePlaintext = encoding.encode(cleartext, params)
+    encoded: types.LwePlaintext = encoding.encode(cleartext, params)  # pyrefly: ignore[not-a-type]
     decoded: types.LweCleartext = encoding.decode(encoded, params)
     self.assertEqual(decoded, cleartext)
 
@@ -63,7 +63,7 @@ def test_encode_decode_succeeds(self, cleartext: types.LweCleartext):
         total_bit_length=32, message_bit_length=16, padding_bit_length=8
     )
 
-    encoded: types.LwePlaintext = encoding.encode(cleartext, params)
+    encoded: types.LwePlaintext = encoding.encode(cleartext, params)  # pyrefly: ignore[not-a-type]
     decoded: types.LweCleartext = encoding.decode(encoded, params)
     self.assertEqual(decoded, cleartext)
 

jaxite/jaxite_cggi/key_switch.py

@@ -16,7 +16,7 @@ class LweKeySwitchingKey:
   """A public key used to switch keys encrypted an LWE ciphertext."""
 
   # the q in Z/qZ, same as the LwePlaintext space
-  modulus: jnp.uint32
+  modulus: jnp.uint32  # pyrefly: ignore[not-a-type]
 
   # the length of the sampled key_data vector s,
   # equal to len(LweCiphertext) - 1.
@@ -65,10 +65,10 @@ def gen_key(
   # The resulting shape of the key switching key is (n, num_levels, lwe_size)
   lwe_ai_samples = prg.uniform(
       shape=(n, num_levels, out_key.lwe_dimension),
-      dtype=jnp.uint32,
+      dtype=jnp.uint32,  # pyrefly: ignore[bad-argument-type]
   )
   lwe_error_samples = prg.rounded_normal(
-      shape=(n, num_levels), dtype=jnp.uint32
+      shape=(n, num_levels), dtype=jnp.uint32  # pyrefly: ignore[bad-argument-type]
   )
   key_data = jnp.zeros((n, num_levels, lwe_size), dtype=jnp.uint32)