pyink reformat

google · Feb 19, 2025 · 9e6d0ce · 9e6d0ce
1 parent c5dab70
commit 9e6d0ce
Show file tree

Hide file tree

Showing 13 changed files with 565 additions and 362 deletions.
diff --git a/frontend/e2e_test.py b/frontend/e2e_test.py
@@ -1,12 +1,13 @@
 from heir import compile
 from heir.mlir import *
 
+
 from absl.testing import absltest  # fmt: skip
 class EndToEndTest(absltest.TestCase):
 
   def test_simple_arithmetic(self):
-    @compile() # defaults to BGV and OpenFHE
-    def foo(a : Secret[I64], b : Secret[I64]):
+    @compile()  # defaults to BGV and OpenFHE
+    def foo(a: Secret[I64], b: Secret[I64]):
       return a * a - b * b
 
     # Test plaintext functionality

diff --git a/frontend/example.py b/frontend/example.py
@@ -5,17 +5,19 @@
 
 # TODO (#1162): Also add the tensorflow-to-tosa-to-HEIR example in example.py, even it doesn't use the main Python frontend?
 
+
 ### Simple Example
-@compile() # defaults to scheme="bgv", OpenFHE backend, and debug=False
-def foo(x : Secret[I16], y : Secret[I16]):
+@compile()  # defaults to scheme="bgv", OpenFHE backend, and debug=False
+def foo(x: Secret[I16], y: Secret[I16]):
   sum = x + y
   diff = x - y
   mul = x * y
   expression = sum * diff + mul
   deadcode = expression * mul
   return expression
 
-foo.setup() # runs keygen/etc
+
+foo.setup()  # runs keygen/etc
 enc_x = foo.encrypt_x(7)
 enc_y = foo.encrypt_y(8)
 result_enc = foo.eval(enc_x, enc_y)
@@ -74,7 +76,6 @@ def foo(x : Secret[I16], y : Secret[I16]):
 # print(f"Expected result for `baz2`: {baz2(7,8,9)}, decrypted result: {result}")
 
 
-
 # ### Custom Pipeline Example
 # @compile(heir_opt_options=["--mlir-to-secret-arithmetic", "--canonicalize", "--cse"], backend=None, debug=True)
 # def foo(x : Secret[I16], y : Secret[I16]):

diff --git a/frontend/heir/backends/dummy.py b/frontend/heir/backends/dummy.py
@@ -2,18 +2,31 @@
 
 from colorama import Fore, Style, init
 
-from heir.core import BackendInterface, CompilationResult,  ClientInterface
+from heir.core import BackendInterface, CompilationResult, ClientInterface
+
 
 class DummyClientInterface(ClientInterface):
 
   def __init__(self, compilation_result: CompilationResult):
     self.compilation_result = compilation_result
 
   def setup(self):
-    print("HEIR Warning (Dummy Backend): " + Fore.YELLOW + Style.BRIGHT + f"{self.compilation_result.func_name}.setup() is a no-op in the Dummy Backend")
+    print(
+        "HEIR Warning (Dummy Backend): "
+        + Fore.YELLOW
+        + Style.BRIGHT
+        + f"{self.compilation_result.func_name}.setup() is a no-op in the Dummy"
+        " Backend"
+    )
 
   def decrypt_result(self, result):
-    print("HEIR Warning (Dummy Backend): " + Fore.YELLOW + Style.BRIGHT + f"{self.compilation_result.func_name}.decrypt() is a no-op in the Dummy Backend")
+    print(
+        "HEIR Warning (Dummy Backend): "
+        + Fore.YELLOW
+        + Style.BRIGHT
+        + f"{self.compilation_result.func_name}.decrypt() is a no-op in the"
+        " Dummy Backend"
+    )
     return result
 
   def __getattr__(self, key):
@@ -22,31 +35,53 @@ def __getattr__(self, key):
       arg_name = key[len("encrypt_") :]
 
       def wrapper(arg):
-        print("HEIR Warning (Dummy Backend): " + Fore.YELLOW + Style.BRIGHT + f"{self.compilation_result.func_name}.{key}() is a no-op in the Dummy Backend")
+        print(
+            "HEIR Warning (Dummy Backend): "
+            + Fore.YELLOW
+            + Style.BRIGHT
+            + f"{self.compilation_result.func_name}.{key}() is a no-op in the"
+            " Dummy Backend"
+        )
         return arg
 
       return wrapper
 
     if key == self.compilation_result.func_name or key == "eval":
-      print("HEIR Warning (Dummy Backend): " + Fore.YELLOW + Style.BRIGHT + f"{self.compilation_result.func_name}.eval() is the same as {self.compilation_result.func_name}() in the Dummy Backend.")
+      print(
+          "HEIR Warning (Dummy Backend): "
+          + Fore.YELLOW
+          + Style.BRIGHT
+          + f"{self.compilation_result.func_name}.eval() is the same as"
+          f" {self.compilation_result.func_name}() in the Dummy Backend."
+      )
       return self.func
 
     raise AttributeError(f"Attribute {key} not found")
 
 
 class DummyBackend(BackendInterface):
 
-    def run_backend(self, workspace_dir, heir_opt, heir_translate, func_name, arg_names, secret_args, heir_opt_output, debug):
-
-        result = CompilationResult(
-              module=None,
-              func_name=func_name,
-              secret_args=secret_args,
-              arg_names=arg_names,
-              arg_enc_funcs=None,
-              result_dec_func=None,
-              main_func=None,
-              setup_funcs=None
-          )
-
-        return DummyClientInterface(result)
+  def run_backend(
+      self,
+      workspace_dir,
+      heir_opt,
+      heir_translate,
+      func_name,
+      arg_names,
+      secret_args,
+      heir_opt_output,
+      debug,
+  ):
+
+    result = CompilationResult(
+        module=None,
+        func_name=func_name,
+        secret_args=secret_args,
+        arg_names=arg_names,
+        arg_enc_funcs=None,
+        result_dec_func=None,
+        main_func=None,
+        setup_funcs=None,
+    )
+
+    return DummyClientInterface(result)
diff --git a/frontend/heir/backends/openfhe/__init__.py b/frontend/heir/backends/openfhe/__init__.py
@@ -1,4 +1,9 @@
 from .backend import OpenFHEBackend
 from .config import OpenFHEConfig, DEFAULT_INSTALLED_OPENFHE_CONFIG, from_os_env
 
-__all__ = ["OpenFHEBackend", "OpenFHEConfig", "DEFAULT_INSTALLED_OPENFHE_CONFIG", "from_os_env"]
+__all__ = [
+    "OpenFHEBackend",
+    "OpenFHEConfig",
+    "DEFAULT_INSTALLED_OPENFHE_CONFIG",
+    "from_os_env",
+]