update example file

heir_py/example.py

@@ -5,41 +5,84 @@
 from heir_py.types import *
 
 # FIXME: Also add the tensorflow-to-tosa-to-HEIR example in here, even it doesn't use Numba
-
-### Simple Example
 # TODO (#1162): Allow manually specifying precision/scale and warn/error if not possible!
 # TODO (#????): Add precision computation/check to Mgmt dialect/infrastructure
-@compile(scheme="ckks", backend="openfhe", debug=True)
-def foo(x : Secret[F32], y : F32):
+# TODO (#1119): expose ctxt serialization in python
+# TODO (#1162) : Fix "ImportError: generic_type: type "PublicKey" is already registered!" when doing setup twice
+
+
+
+
+### Simple Example
+@compile() # defaults to scheme="bgv"", backend="openfhe", debug=False
+def foo(x : Secret[I16], y : Secret[I16]):
   sum = x + y
-  diff = x + y
+  diff = x - y
   mul = x * y
   expression = sum * diff + mul
   deadcode = expression * mul
   return expression
 
 foo.setup() # runs keygen/etc
-enc_x = foo.encrypt_x(7) # TODO(#1119): expose ctxt serialization in python
-# enc_y = foo.encrypt_y(8)
-result_enc = foo.eval(enc_x, 8)
+enc_x = foo.encrypt_x(7)
+enc_y = foo.encrypt_y(8)
+result_enc = foo.eval(enc_x, enc_y)
 result = foo.decrypt_result(result_enc)
-print(f"Expected result: {foo(7,8)}, decrypted result: {result}")
+print(f"Expected result for `foo`: {foo(7,8)}, decrypted result: {result}")
+
+
+
 
+# ### CKKS Example
+# @compile(scheme="ckks") # other options default to backend="openfhe", debug=False
+# def bar(x : Secret[F32], y : Secret[F32]):
+#   sum = x + y
+#   diff = x - y
+#   mul = x * y
+#   expression = sum * diff + mul
+#   deadcode = expression * mul
+#   return expression
 
-# # Matmul Example
-# @compile("float32[:,:](float32[:,:],float32[:,:])", backend="openfhe", debug=True)
-# def goo(a : Secret[Tensor[4,4,F32]], b : Secret[Tensor[4,4,F32]]):
+# bar.setup() # runs keygen/etc
+# enc_x = bar.encrypt_x(7)
+# enc_y = bar.encrypt_y(8)
+# result_enc = bar.eval(enc_x, enc_y)
+# result = bar.decrypt_result(result_enc)
+# print(f"Expected result for `bar`: {bar(7,8)}, decrypted result: {result}")
+
+
+
+
+# ### Ciphertext-Plaintext Example
+# @compile()
+# def baz(x: Secret[I16], y : I16):
+#   ptxt_mul = x * y
+#   ctxt_mul = x * x
+#   return ptxt_mul + ctxt_mul
+
+# baz.setup() # runs keygen/etc
+# enc_x = baz.encrypt_x(7)
+# result_enc = baz.eval(enc_x, 8)
+# result = baz.decrypt_result(result_enc)
+# print(f"Expected result for `baz`: {baz(7,8)}, decrypted result: {result}")
+
+
+
+
+# ### Matmul Example
+# @compile(scheme='ckks')
+# def qux(a : Secret[Tensor[4,4,F32]], b : Secret[Tensor[4,4,F32]]):
 #  AB = matmul(a,b)
 #  AABB = matmul(a+a, b+b)
 #  return AB + AABB
 
 # a = np.array([[1,2],[3,4]])
 # b = np.array([[5,6],[7,8]])
-# print(goo(a,b))
-
-# goo.setup()
-# enc_a = goo.encrypt_a(a)
-# enc_b = goo.encrypt_b(b)
-# result_enc = goo.eval(enc_a, enc_b)
-# result = goo.decrypt_result(result_enc)
-# print(f"Expected result: {np.matmul(a,b)}, decrypted result: {result}")
+# print(qux(a,b))
+
+# qux.setup()
+# enc_a = qux.encrypt_a(a)
+# enc_b = qux.encrypt_b(b)
+# result_enc = qux.eval(enc_a, enc_b)
+# result = qux.decrypt_result(result_enc)
+# print(f"Expected result: {qux(a,b)}, decrypted result: {result}")