[WIP] More elaborate Paillier example

examples/paillier.py

@@ -1,26 +1,73 @@
+import numpy as np
 import tensorflow as tf
 import tf_big
 
+# use secure operations by default
 tf_big.set_secure_default(True)
 
-p = tf_big.constant([17])
-q = tf_big.constant([19])
-n = p * q
+class EncryptionKey:
+  def __init__(self, n, nn, g):
+    self.n = n
+    self.nn = nn
+    self.g = g
 
-g = n + 1
-nn = n * n
+class DecryptionKey:
+  def __init__(self):
+    # TODO what do we need?
+    pass
 
-x = tf.constant([[4]])
-# r = tf_big.random.uniform(n)
-r = tf_big.constant([82])
-assert r.shape == x.shape, (r.shape, x.shape)
+def dummy_keygen():
+  # TODO use fixed large primes  
+  p = tf_big.constant([17])
+  q = tf_big.constant([19])
+  n = p * q
 
-gx = tf_big.pow(g, x, nn, secure=True)
-assert gx.shape.as_list() == [1, 1], gx.shape
-rn = tf_big.pow(r, n, nn, secure=True)
-assert rn.shape.as_list() == [1, 1], rn.shape
-c = gx * rn #% nn
+  g = n + 1
+  nn = n * n
+
+  ek = EncryptionKey(n, nn, g)
+  dk = DecryptionKey()
+  return ek, dk
+
+def encrypt(ek, x):
+  # r = tf_big.random.uniform(ek.n, shape=x.shape)
+  r = tf_big.convert_to_tensor(tf.constant([[123, 124], [125, 126]])) # TODO
+  assert r.shape == x.shape, "Shapes are not matching: {}, {}".format(r.shape, x.shape)
+
+  gx = tf_big.pow(ek.g, x, ek.nn)
+  rn = tf_big.pow(r, ek.n, ek.nn)
+  c = gx * rn % ek.nn
+  return c
+
+def decrypt(dk, c):
+  # TODO what do we need?
+  pass
+
+def add(ek, c1, c2):
+  c = c1 * c2 % ek.nn
+  return c
+
+def mul(ek, c1, x2):
+  c = tf_big.pow(c1, x2, ek.nn)
+  return c
+
+ek, dk = dummy_keygen()
+
+# TODO(Morten) relace with lin reg computation?
+
+x1 = tf.constant([[1, 2], [3, 4]])
+c1 = encrypt(ek, x1)
+
+x1 = tf.constant([[5, 6], [7, 8]])
+c2 = encrypt(ek, x2)
+
+c3 = add(ek, c1, c2)
+
+c4 = mul(ek, c3, tf.constant(3))
+
+y = decrypt(dk, c4)
 
 with tf.Session() as sess:
-  res = sess.run(c)
-  print(res)
+  actual = sess.run(y)
+  expected = (x1 + x2) * 3
+  np.testing.assert_array_equal(actual, expected)